[FLINK-38267] Only call channel state rescaling logic for exchange with channel state to avoid UnsupportedOperationException #26931

1996fanrui · 2025-08-20T12:53:07Z

What is the purpose of the change

Job cannot be recovered from UC(unaligned checkpoint) after rescaling, and the exception is:

java.lang.UnsupportedOperationException: Cannot rescale the given pointwise partitioner. 
Did you change the partitioner to forward or rescale? 
It may also help to add an explicit shuffle().

When one upstream task has multiple output exchanges, which including UC SUPPORTED exchanges(likes hash or rebalance) and at least one UC UNSUPPORTED exchanges(likes Forward or rescale).

Or when one downstream task has multiple input exchanges, which including UC SUPPORTED exchanges(likes hash or rebalance) and at least one UC UNSUPPORTED exchanges(likes Forward or rescale).

See more from FLINK-38267

Brief change log

The implemented solution was to make the state redistribution logic more granular by checking for in-flight data on a per-exchange basis instead of a per-task basis.

[FLINK-38267][checkpoint] Refactor hasInputState and hasOutputState related logic in TaskStateAssignment
- Precise State Tracking: The TaskStateAssignment class was refactored to no longer use a simple boolean flag. It now precisely tracks which specific input gates and result partitions contain in-flight data.
[FLINK-38267] Only call channel state rescaling logic for exchange with channel state to avoid UnsupportedOperationException
- Per-Channel/Partition Checks: The core redistribution methods, reDistributeInputChannelStates and reDistributeResultSubpartitionStates, were modified. Their internal logic now iterates through each input gate or output partition and uses new helper methods (hasInFlightDataForInputGate and hasInFlightDataForResultPartition) to check if that specific channel has state.
- Conditional Logic: The state redistribution logic is now wrapped in a conditional block. It is only invoked for a channel if the per-exchange check passes. This ensures that stateless exchanges (like forward or rescale) are correctly skipped, avoiding the exception.

This approach fixes the bug by applying the redistribution logic only where it is actually needed, allowing jobs with mixed partitioner types to rescale from an unaligned checkpoint successfully.

Verifying this change

Doing

Does this pull request potentially affect one of the following parts:

Dependencies (does it add or upgrade a dependency): no
The public API, i.e., is any changed class annotated with @Public(Evolving): no
The serializers: no
The runtime per-record code paths (performance sensitive): no
Anything that affects deployment or recovery: JobManager (and its components), Checkpointing, Kubernetes/Yarn, ZooKeeper: no
The S3 file system connector: no

Documentation

Does this pull request introduce a new feature? no

…elated logic in TaskStateAssignment

flinkbot · 2025-08-20T12:59:26Z

CI report:

6d7bec8 Azure: FAILURE
2e58646 Azure: PENDING

Bot commands

The @flinkbot bot supports the following commands:

@flinkbot run azure re-run the last Azure build

gaborgsomogyi

Going to do some manual testing in the meantime

gaborgsomogyi · 2025-08-21T10:56:50Z

...ntime/src/main/java/org/apache/flink/runtime/checkpoint/InflightDataRescalingDescriptor.java

+                new InflightDataGateOrPartitionRescalingDescriptor(
+                        new int[0],
+                        RescaleMappings.identity(0, 0),
+                        java.util.Collections.emptySet(),


Nit: Maybe it can be imported.

Good catch! And updated

gaborgsomogyi · 2025-08-21T11:23:45Z

flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/TaskStateAssignment.java

+                    Arrays.asList(inputResult.getProducer().getProducedDataSets())
+                            .indexOf(inputResult);


Are we sure that this is working all the time because IntermediateResult has no valid equals implementation and just uses == on the references?

From the current code implementation, it works well since we do not create multiple IntermediateResult for same object.

But I think it is not a good way to find the partitionIndex, I have updated the code that find partitionIndex via the IntermediateDataSetID.

gaborgsomogyi · 2025-08-21T11:24:55Z

flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/TaskStateAssignment.java

+            int gateIndex =
+                    downstreamAssignment.executionJobVertex.getInputs().indexOf(producedResult);


gaborgsomogyi · 2025-08-21T12:27:10Z

.../org/apache/flink/test/checkpointing/UnalignedCheckpointRescaleWithMixedExchangesITCase.java

+ * an unaligned checkpoint.
+ */
+@RunWith(Parameterized.class)
+public class UnalignedCheckpointRescaleWithMixedExchangesITCase extends TestLogger {


I've gone through the existing and these new UC tests which are providing excellent coverage. What I've not yet found is cases which skip rescaling when an edge has no in-flight/channel state.

This ITCase has 4 DAGs, and all of them have forward or rescale exchange.

These 2 exchanges does not support unaligned checkpoint, resulting in the corresponding edge has no in-flight/channel state.

gaborgsomogyi · 2025-08-21T12:34:43Z

.../org/apache/flink/test/checkpointing/UnalignedCheckpointRescaleWithMixedExchangesITCase.java

+        // Decrease the memory segment size to avoid the test is so slow for some reasons:
+        // 1. Recovery phase needs to consume all inflight buffers
+        // 2. Forward or rescale exchange does not support unaligned checkpoint.
+        conf.set(TaskManagerOptions.MEMORY_SEGMENT_SIZE, MemorySize.parse("1 kb"));


Not yet understand these comments, could you elaborate on this? On my side if I remove it then it makes tests just flaky.

Not yet understand these comments, could you elaborate on this?

I added more comments, I hope it is clear.

On my side if I remove it then it makes tests just flaky.

Sorry, do you mean is it slow or failed on your local? As I understand, it get slower, but should not be failed.

If I remove the mentioned line then majority of the time the last test fails, but in general it's flaky🤷🏻‍♂️

I needed to execute the class 4-5 times to see the following failure:

java.lang.IllegalStateException: Job has entered FAILED state, but expecting [RUNNING]

I try to run this ITCase 10 times on my local after removing conf.set(TaskManagerOptions.MEMORY_SEGMENT_SIZE, MemorySize.parse("1 kb"));, it works well.

Could you please mvn install this PR in your local, and try run test again. If the flink-runtime code is master branch, this ITCase will be failed (It is exactly why I introduce this ITCase, it could reproduce the FLINK-38267 bug). I encountered this issue during development as well, and it works well after installing this PR to local.

If it is still failed after re-installing, would you mind providing the detailed exception info and exception stack? why Job has entered FAILED state? Thank you very much!

gaborgsomogyi · 2025-08-21T12:47:52Z

flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/TaskStateAssignment.java

@@ -204,15 +266,16 @@ public OperatorSubtaskState getSubtaskState(OperatorInstanceID instanceID) {
                                    return assignment.getInputMapping(assignmentIndex, recompute);
                                },
                                outputSubtaskMappings,
-                                this::getOutputMapping))
+                                this::getOutputMapping,
+                                false))
                .build();
    }

    public boolean hasUpstreamOutputStates() {
        if (hasUpstreamOutputStates == null) {
            hasUpstreamOutputStates =
                    Arrays.stream(getUpstreamAssignments())


This issue/risk is not introduced here but maybe we could add a check while we're hanging around here. Namely getUpstreamAssignments is depending on executionJobVertex.getInputs() which is nullable.

Sorry, I do not completely understand this comment.

Do you mean to add checkState(executionJobVertex.getInputs() != null) or something like that here?

Yeah, I think that's enough to be on the safe side.

ExecutionJobVertex#inputs is @Nullable, but ExecutionJobVertex#getInputs is not sincegetInputs() checked the isInitialized(), and inputs never be null once is isInitialized.

public List<IntermediateResult> getInputs() { checkState(isInitialized()); return inputs; }

WDYT?

1996fanrui

Thanks for the review, all comments are addressed.

1996fanrui · 2025-08-21T14:01:27Z

...ntime/src/main/java/org/apache/flink/runtime/checkpoint/InflightDataRescalingDescriptor.java

+                new InflightDataGateOrPartitionRescalingDescriptor(
+                        new int[0],
+                        RescaleMappings.identity(0, 0),
+                        java.util.Collections.emptySet(),


Good catch! And updated

1996fanrui · 2025-08-21T14:12:17Z

flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/TaskStateAssignment.java

+                    Arrays.asList(inputResult.getProducer().getProducedDataSets())
+                            .indexOf(inputResult);


From the current code implementation, it works well since we do not create multiple IntermediateResult for same object.

But I think it is not a good way to find the partitionIndex, I have updated the code that find partitionIndex via the IntermediateDataSetID.

1996fanrui · 2025-08-21T14:12:24Z

flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/TaskStateAssignment.java

+            int gateIndex =
+                    downstreamAssignment.executionJobVertex.getInputs().indexOf(producedResult);


1996fanrui · 2025-08-21T14:15:25Z

.../org/apache/flink/test/checkpointing/UnalignedCheckpointRescaleWithMixedExchangesITCase.java

+ * an unaligned checkpoint.
+ */
+@RunWith(Parameterized.class)
+public class UnalignedCheckpointRescaleWithMixedExchangesITCase extends TestLogger {


This ITCase has 4 DAGs, and all of them have forward or rescale exchange.

These 2 exchanges does not support unaligned checkpoint, resulting in the corresponding edge has no in-flight/channel state.

1996fanrui · 2025-08-21T14:36:41Z

.../org/apache/flink/test/checkpointing/UnalignedCheckpointRescaleWithMixedExchangesITCase.java

+        // Decrease the memory segment size to avoid the test is so slow for some reasons:
+        // 1. Recovery phase needs to consume all inflight buffers
+        // 2. Forward or rescale exchange does not support unaligned checkpoint.
+        conf.set(TaskManagerOptions.MEMORY_SEGMENT_SIZE, MemorySize.parse("1 kb"));


Not yet understand these comments, could you elaborate on this?

I added more comments, I hope it is clear.

On my side if I remove it then it makes tests just flaky.

Sorry, do you mean is it slow or failed on your local? As I understand, it get slower, but should not be failed.

1996fanrui · 2025-08-21T14:47:43Z

flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/TaskStateAssignment.java

@@ -204,15 +266,16 @@ public OperatorSubtaskState getSubtaskState(OperatorInstanceID instanceID) {
                                    return assignment.getInputMapping(assignmentIndex, recompute);
                                },
                                outputSubtaskMappings,
-                                this::getOutputMapping))
+                                this::getOutputMapping,
+                                false))
                .build();
    }

    public boolean hasUpstreamOutputStates() {
        if (hasUpstreamOutputStates == null) {
            hasUpstreamOutputStates =
                    Arrays.stream(getUpstreamAssignments())


Sorry, I do not completely understand this comment.

Do you mean to add checkState(executionJobVertex.getInputs() != null) or something like that here?

davidradl · 2025-08-21T15:07:38Z

flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/TaskStateAssignment.java

+                                        .getInfos().stream()
+                                                .mapToInt(ResultSubpartitionInfo::getPartitionIdx);
+                            } else {
+                                throw new IllegalStateException(


Can this be driven? If so please could you add a test to drive this throw?

There is no uncovered implementation so can't be tested now.

Yes, this is a defense mechanism to warn developers if other subclasses are introduced in the future.

Unless we introduce a test subclass for this testing.

… exchange with channel state to avoid UnsupportedOperationException

[FLINK-38267][checkpoint] Refactor hasInputState and hasOutputState r…

c33e974

…elated logic in TaskStateAssignment

1996fanrui force-pushed the 38267/fix-UC-multi-exchanges branch 2 times, most recently from 1e6b2e7 to 4f49de1 Compare August 20, 2025 20:39

gaborgsomogyi reviewed Aug 21, 2025

View reviewed changes

1996fanrui marked this pull request as ready for review August 21, 2025 13:40

1996fanrui commented Aug 21, 2025

View reviewed changes

davidradl reviewed Aug 21, 2025

View reviewed changes

github-actions bot added the community-reviewed PR has been reviewed by the community. label Aug 21, 2025

[FLINK-38267][checkpoint] Only call channel state rescaling logic for…

2e58646

… exchange with channel state to avoid UnsupportedOperationException

1996fanrui force-pushed the 38267/fix-UC-multi-exchanges branch from 6d7bec8 to 2e58646 Compare August 21, 2025 18:52

		Arrays.asList(inputResult.getProducer().getProducedDataSets())
		.indexOf(inputResult);

		int gateIndex =
		downstreamAssignment.executionJobVertex.getInputs().indexOf(producedResult);

[FLINK-38267] Only call channel state rescaling logic for exchange with channel state to avoid UnsupportedOperationException #26931

Are you sure you want to change the base?

[FLINK-38267] Only call channel state rescaling logic for exchange with channel state to avoid UnsupportedOperationException #26931

Conversation

1996fanrui commented Aug 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What is the purpose of the change

Brief change log

Verifying this change

Does this pull request potentially affect one of the following parts:

Documentation

Uh oh!

flinkbot commented Aug 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

CI report:

Uh oh!

gaborgsomogyi left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

1996fanrui Aug 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

1996fanrui left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

1996fanrui commented Aug 20, 2025 •

edited

Loading

flinkbot commented Aug 20, 2025 •

edited

Loading

1996fanrui Aug 21, 2025 •

edited

Loading