This is an automated email from the ASF dual-hosted git repository. 1996fanrui pushed a commit to branch release-1.20 in repository https://gitbox.apache.org/repos/asf/flink.git
commit 3f1abf2fff4b3f1228a46c6f980d9a4eadce6a9d 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 | 4 + .../org/apache/flink/runtime/jobgraph/JobEdge.java | 16 ++- .../scheduler/adaptive/AdaptiveScheduler.java | 5 +- .../scheduler/adaptive/ForwardEdgesAdapter.java | 71 ++++++++++ .../adaptive/ForwardEdgesAdapterTest.java | 151 +++++++++++++++++++++ 5 files changed, 245 insertions(+), 2 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 2aad2613996..972e5b3dcf8 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 @@ -103,6 +103,10 @@ public class IntermediateDataSet implements java.io.Serializable { consumers.add(edge); } + public void updateDistributionPattern(DistributionPattern distributionPattern) { + this.distributionPattern = distributionPattern; + } + // -------------------------------------------------------------------------------------------- @Override 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 4649303c144..bc2c941df90 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; @@ -87,6 +87,20 @@ public class JobEdge implements java.io.Serializable { this.isBroadcast = isBroadcast; } + 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(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 e5bc372f22d..d95a8e62960 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 @@ -150,6 +150,7 @@ import java.util.function.Function; import java.util.stream.Collectors; 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 @@ -1261,8 +1262,10 @@ public class AdaptiveScheduler @Nonnull private ExecutionGraph createExecutionGraphAndRestoreState( VertexParallelismStore adjustedParallelismStore) throws Exception { + 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..d05e43956db --- /dev/null +++ b/flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/adaptive/ForwardEdgesAdapter.java @@ -0,0 +1,71 @@ +/* + * 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.io.network.api.writer.SubtaskStateMapper; +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 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()) { + final boolean forwardable = isSameParallelism(input, vertexParallelismStore); + input.updateDistributionPattern( + forwardable, + forwardable ? POINTWISE : ALL_TO_ALL, + forwardable + ? SubtaskStateMapper.UNSUPPORTED + : SubtaskStateMapper.ARBITRARY, + forwardable + ? SubtaskStateMapper.UNSUPPORTED + : SubtaskStateMapper.ROUND_ROBIN, + forwardable ? "FORWARD" : "REBALANCE"); + } + } + } + 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..eb11500a6a6 --- /dev/null +++ b/flink-runtime/src/test/java/org/apache/flink/runtime/scheduler/adaptive/ForwardEdgesAdapterTest.java @@ -0,0 +1,151 @@ +/* + * 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.io.network.api.writer.SubtaskStateMapper; +import org.apache.flink.runtime.io.network.partition.ResultPartitionType; +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.JobGraphTestUtils; +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.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() { + JobVertex source = new JobVertex("source"); + JobVertex middle = new JobVertex("middle"); + JobVertex firstSink = new JobVertex("firstSink"); + JobVertex secondSink = new JobVertex("secondSink"); + for (JobVertex vertex : new JobVertex[] {source, middle, firstSink, secondSink}) { + vertex.setParallelism(2); + } + middle.connectNewDataSetAsInput(source, POINTWISE, ResultPartitionType.PIPELINED) + .setForward(true); + firstSink + .connectNewDataSetAsInput(middle, POINTWISE, ResultPartitionType.PIPELINED) + .setForward(true); + secondSink + .connectNewDataSetAsInput(middle, POINTWISE, ResultPartitionType.PIPELINED) + .setForward(true); + return JobGraphTestUtils.streamingJobGraph(source, middle, firstSink, secondSink); + } + + 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); + assertThat(adjustedInput.isForward()).isEqualTo(forwardable); + assertThat(adjustedInput.getDistributionPattern()) + .isEqualTo(forwardable ? POINTWISE : ALL_TO_ALL); + assertThat(adjustedInput.getUpstreamSubtaskStateMapper()) + .isEqualTo( + forwardable + ? SubtaskStateMapper.UNSUPPORTED + : SubtaskStateMapper.ARBITRARY); + assertThat(adjustedInput.getDownstreamSubtaskStateMapper()) + .isEqualTo( + forwardable + ? SubtaskStateMapper.UNSUPPORTED + : SubtaskStateMapper.ROUND_ROBIN); + assertThat(adjustedInput.getShipStrategyName()) + .isEqualTo(forwardable ? "FORWARD" : "REBALANCE"); + + IntermediateDataSet adjustedInputSource = adjustedInput.getSource(); + assertThat(adjustedInputSource.getDistributionPattern()) + .isEqualTo(forwardable ? POINTWISE : ALL_TO_ALL); + } + } + } + } +}
