| /* |
| * 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.checkpoint; |
| |
| import org.apache.flink.annotation.Internal; |
| import org.apache.flink.annotation.VisibleForTesting; |
| import org.apache.flink.api.java.tuple.Tuple2; |
| import org.apache.flink.runtime.OperatorIDPair; |
| import org.apache.flink.runtime.checkpoint.channel.InputChannelInfo; |
| import org.apache.flink.runtime.checkpoint.channel.ResultSubpartitionInfo; |
| import org.apache.flink.runtime.executiongraph.Execution; |
| import org.apache.flink.runtime.executiongraph.ExecutionJobVertex; |
| import org.apache.flink.runtime.executiongraph.IntermediateResult; |
| import org.apache.flink.runtime.io.network.api.writer.SubtaskStateMapper; |
| import org.apache.flink.runtime.jobgraph.IntermediateDataSet; |
| import org.apache.flink.runtime.jobgraph.IntermediateDataSetID; |
| import org.apache.flink.runtime.jobgraph.JobEdge; |
| import org.apache.flink.runtime.jobgraph.OperatorID; |
| import org.apache.flink.runtime.jobgraph.OperatorInstanceID; |
| import org.apache.flink.runtime.state.AbstractChannelStateHandle; |
| import org.apache.flink.runtime.state.InputChannelStateHandle; |
| import org.apache.flink.runtime.state.KeyGroupRange; |
| import org.apache.flink.runtime.state.KeyGroupRangeAssignment; |
| import org.apache.flink.runtime.state.KeyGroupsStateHandle; |
| import org.apache.flink.runtime.state.KeyedStateHandle; |
| import org.apache.flink.runtime.state.ResultSubpartitionStateHandle; |
| import org.apache.flink.runtime.state.StateObject; |
| import org.apache.flink.util.CollectionUtil; |
| import org.apache.flink.util.Preconditions; |
| |
| import org.slf4j.Logger; |
| import org.slf4j.LoggerFactory; |
| |
| import java.util.ArrayList; |
| import java.util.Collection; |
| import java.util.HashMap; |
| import java.util.HashSet; |
| import java.util.List; |
| import java.util.Map; |
| import java.util.Set; |
| import java.util.function.Function; |
| import java.util.stream.Collectors; |
| |
| import static java.util.Collections.emptyList; |
| import static org.apache.flink.util.Preconditions.checkNotNull; |
| |
| /** |
| * This class encapsulates the operation of assigning restored state when restoring from a |
| * checkpoint. |
| */ |
| @Internal |
| public class StateAssignmentOperation { |
| |
| private static final Logger LOG = LoggerFactory.getLogger(StateAssignmentOperation.class); |
| |
| private final Set<ExecutionJobVertex> tasks; |
| private final Map<OperatorID, OperatorState> operatorStates; |
| |
| private final long restoreCheckpointId; |
| private final boolean allowNonRestoredState; |
| |
| /** The state assignments for each ExecutionJobVertex that will be filled in multiple passes. */ |
| private final Map<ExecutionJobVertex, TaskStateAssignment> vertexAssignments; |
| /** |
| * Stores the assignment of a consumer. {@link IntermediateResult} only allows to traverse |
| * producer. |
| */ |
| private final Map<IntermediateDataSetID, TaskStateAssignment> consumerAssignment = |
| new HashMap<>(); |
| |
| public StateAssignmentOperation( |
| long restoreCheckpointId, |
| Set<ExecutionJobVertex> tasks, |
| Map<OperatorID, OperatorState> operatorStates, |
| boolean allowNonRestoredState) { |
| |
| this.restoreCheckpointId = restoreCheckpointId; |
| this.tasks = Preconditions.checkNotNull(tasks); |
| this.operatorStates = Preconditions.checkNotNull(operatorStates); |
| this.allowNonRestoredState = allowNonRestoredState; |
| this.vertexAssignments = CollectionUtil.newHashMapWithExpectedSize(tasks.size()); |
| } |
| |
| public void assignStates() { |
| checkStateMappingCompleteness(allowNonRestoredState, operatorStates, tasks); |
| |
| Map<OperatorID, OperatorState> localOperators = new HashMap<>(operatorStates); |
| |
| // find the states of all operators belonging to this task and compute additional |
| // information in first pass |
| for (ExecutionJobVertex executionJobVertex : tasks) { |
| List<OperatorIDPair> operatorIDPairs = executionJobVertex.getOperatorIDs(); |
| Map<OperatorID, OperatorState> operatorStates = |
| CollectionUtil.newHashMapWithExpectedSize(operatorIDPairs.size()); |
| for (OperatorIDPair operatorIDPair : operatorIDPairs) { |
| OperatorID operatorID = |
| operatorIDPair |
| .getUserDefinedOperatorID() |
| .filter(localOperators::containsKey) |
| .orElse(operatorIDPair.getGeneratedOperatorID()); |
| |
| OperatorState operatorState = localOperators.remove(operatorID); |
| if (operatorState == null) { |
| operatorState = |
| new OperatorState( |
| operatorID, |
| executionJobVertex.getParallelism(), |
| executionJobVertex.getMaxParallelism()); |
| } |
| operatorStates.put(operatorIDPair.getGeneratedOperatorID(), operatorState); |
| } |
| |
| final TaskStateAssignment stateAssignment = |
| new TaskStateAssignment( |
| executionJobVertex, |
| operatorStates, |
| consumerAssignment, |
| vertexAssignments); |
| vertexAssignments.put(executionJobVertex, stateAssignment); |
| for (final IntermediateResult producedDataSet : executionJobVertex.getInputs()) { |
| consumerAssignment.put(producedDataSet.getId(), stateAssignment); |
| } |
| } |
| |
| // repartition state |
| for (TaskStateAssignment stateAssignment : vertexAssignments.values()) { |
| if (stateAssignment.hasNonFinishedState |
| // FLINK-31963: We need to run repartitioning for stateless operators that have |
| // upstream output or downstream input states. |
| || stateAssignment.hasUpstreamOutputStates() |
| || stateAssignment.hasDownstreamInputStates()) { |
| assignAttemptState(stateAssignment); |
| } |
| } |
| |
| // actually assign the state |
| for (TaskStateAssignment stateAssignment : vertexAssignments.values()) { |
| // If upstream has output states or downstream has input states, even the empty task |
| // state should be assigned for the current task in order to notify this task that the |
| // old states will send to it which likely should be filtered. |
| if (stateAssignment.hasNonFinishedState |
| || stateAssignment.isFullyFinished |
| || stateAssignment.hasUpstreamOutputStates() |
| || stateAssignment.hasDownstreamInputStates()) { |
| assignTaskStateToExecutionJobVertices(stateAssignment); |
| } |
| } |
| } |
| |
| private void assignAttemptState(TaskStateAssignment taskStateAssignment) { |
| |
| // 1. first compute the new parallelism |
| checkParallelismPreconditions(taskStateAssignment); |
| |
| List<KeyGroupRange> keyGroupPartitions = |
| createKeyGroupPartitions( |
| taskStateAssignment.executionJobVertex.getMaxParallelism(), |
| taskStateAssignment.newParallelism); |
| |
| /* |
| * Redistribute ManagedOperatorStates and RawOperatorStates from old parallelism to new parallelism. |
| * |
| * The old ManagedOperatorStates with old parallelism 3: |
| * |
| * parallelism0 parallelism1 parallelism2 |
| * op0 states0,0 state0,1 state0,2 |
| * op1 |
| * op2 states2,0 state2,1 state2,2 |
| * op3 states3,0 state3,1 state3,2 |
| * |
| * The new ManagedOperatorStates with new parallelism 4: |
| * |
| * parallelism0 parallelism1 parallelism2 parallelism3 |
| * op0 state0,0 state0,1 state0,2 state0,3 |
| * op1 |
| * op2 state2,0 state2,1 state2,2 state2,3 |
| * op3 state3,0 state3,1 state3,2 state3,3 |
| */ |
| reDistributePartitionableStates( |
| taskStateAssignment.oldState, |
| taskStateAssignment.newParallelism, |
| OperatorSubtaskState::getManagedOperatorState, |
| RoundRobinOperatorStateRepartitioner.INSTANCE, |
| taskStateAssignment.subManagedOperatorState); |
| reDistributePartitionableStates( |
| taskStateAssignment.oldState, |
| taskStateAssignment.newParallelism, |
| OperatorSubtaskState::getRawOperatorState, |
| RoundRobinOperatorStateRepartitioner.INSTANCE, |
| taskStateAssignment.subRawOperatorState); |
| |
| reDistributeInputChannelStates(taskStateAssignment); |
| reDistributeResultSubpartitionStates(taskStateAssignment); |
| |
| reDistributeKeyedStates(keyGroupPartitions, taskStateAssignment); |
| } |
| |
| private void assignTaskStateToExecutionJobVertices(TaskStateAssignment assignment) { |
| ExecutionJobVertex executionJobVertex = assignment.executionJobVertex; |
| |
| List<OperatorIDPair> operatorIDs = executionJobVertex.getOperatorIDs(); |
| final int newParallelism = executionJobVertex.getParallelism(); |
| |
| /* |
| * An executionJobVertex's all state handles needed to restore are something like a matrix |
| * |
| * parallelism0 parallelism1 parallelism2 parallelism3 |
| * op0 sh(0,0) sh(0,1) sh(0,2) sh(0,3) |
| * op1 sh(1,0) sh(1,1) sh(1,2) sh(1,3) |
| * op2 sh(2,0) sh(2,1) sh(2,2) sh(2,3) |
| * op3 sh(3,0) sh(3,1) sh(3,2) sh(3,3) |
| * |
| */ |
| for (int subTaskIndex = 0; subTaskIndex < newParallelism; subTaskIndex++) { |
| Execution currentExecutionAttempt = |
| executionJobVertex.getTaskVertices()[subTaskIndex].getCurrentExecutionAttempt(); |
| |
| if (assignment.isFullyFinished) { |
| assignFinishedStateToTask(currentExecutionAttempt); |
| } else { |
| assignNonFinishedStateToTask( |
| assignment, operatorIDs, subTaskIndex, currentExecutionAttempt); |
| } |
| } |
| } |
| |
| private void assignFinishedStateToTask(Execution currentExecutionAttempt) { |
| JobManagerTaskRestore taskRestore = |
| new JobManagerTaskRestore( |
| restoreCheckpointId, TaskStateSnapshot.FINISHED_ON_RESTORE); |
| currentExecutionAttempt.setInitialState(taskRestore); |
| } |
| |
| private void assignNonFinishedStateToTask( |
| TaskStateAssignment assignment, |
| List<OperatorIDPair> operatorIDs, |
| int subTaskIndex, |
| Execution currentExecutionAttempt) { |
| TaskStateSnapshot taskState = new TaskStateSnapshot(operatorIDs.size(), false); |
| |
| for (OperatorIDPair operatorID : operatorIDs) { |
| OperatorInstanceID instanceID = |
| OperatorInstanceID.of(subTaskIndex, operatorID.getGeneratedOperatorID()); |
| |
| OperatorSubtaskState operatorSubtaskState = assignment.getSubtaskState(instanceID); |
| |
| taskState.putSubtaskStateByOperatorID( |
| operatorID.getGeneratedOperatorID(), operatorSubtaskState); |
| } |
| |
| JobManagerTaskRestore taskRestore = |
| new JobManagerTaskRestore(restoreCheckpointId, taskState); |
| currentExecutionAttempt.setInitialState(taskRestore); |
| } |
| |
| public void checkParallelismPreconditions(TaskStateAssignment taskStateAssignment) { |
| for (OperatorState operatorState : taskStateAssignment.oldState.values()) { |
| checkParallelismPreconditions(operatorState, taskStateAssignment.executionJobVertex); |
| } |
| } |
| |
| private void reDistributeKeyedStates( |
| List<KeyGroupRange> keyGroupPartitions, TaskStateAssignment stateAssignment) { |
| stateAssignment.oldState.forEach( |
| (operatorID, operatorState) -> { |
| for (int subTaskIndex = 0; |
| subTaskIndex < stateAssignment.newParallelism; |
| subTaskIndex++) { |
| OperatorInstanceID instanceID = |
| OperatorInstanceID.of(subTaskIndex, operatorID); |
| Tuple2<List<KeyedStateHandle>, List<KeyedStateHandle>> subKeyedStates = |
| reAssignSubKeyedStates( |
| operatorState, |
| keyGroupPartitions, |
| subTaskIndex, |
| stateAssignment.newParallelism, |
| operatorState.getParallelism()); |
| stateAssignment.subManagedKeyedState.put(instanceID, subKeyedStates.f0); |
| stateAssignment.subRawKeyedState.put(instanceID, subKeyedStates.f1); |
| } |
| }); |
| } |
| |
| // TODO rewrite based on operator id |
| private Tuple2<List<KeyedStateHandle>, List<KeyedStateHandle>> reAssignSubKeyedStates( |
| OperatorState operatorState, |
| List<KeyGroupRange> keyGroupPartitions, |
| int subTaskIndex, |
| int newParallelism, |
| int oldParallelism) { |
| |
| List<KeyedStateHandle> subManagedKeyedState; |
| List<KeyedStateHandle> subRawKeyedState; |
| |
| if (newParallelism == oldParallelism) { |
| if (operatorState.getState(subTaskIndex) != null) { |
| subManagedKeyedState = |
| operatorState.getState(subTaskIndex).getManagedKeyedState().asList(); |
| subRawKeyedState = operatorState.getState(subTaskIndex).getRawKeyedState().asList(); |
| } else { |
| subManagedKeyedState = emptyList(); |
| subRawKeyedState = emptyList(); |
| } |
| } else { |
| subManagedKeyedState = |
| getManagedKeyedStateHandles( |
| operatorState, keyGroupPartitions.get(subTaskIndex)); |
| subRawKeyedState = |
| getRawKeyedStateHandles(operatorState, keyGroupPartitions.get(subTaskIndex)); |
| } |
| |
| if (subManagedKeyedState.isEmpty() && subRawKeyedState.isEmpty()) { |
| return new Tuple2<>(emptyList(), emptyList()); |
| } else { |
| return new Tuple2<>(subManagedKeyedState, subRawKeyedState); |
| } |
| } |
| |
| public static <T extends StateObject> void reDistributePartitionableStates( |
| Map<OperatorID, OperatorState> oldOperatorStates, |
| int newParallelism, |
| Function<OperatorSubtaskState, StateObjectCollection<T>> extractHandle, |
| OperatorStateRepartitioner<T> stateRepartitioner, |
| Map<OperatorInstanceID, List<T>> result) { |
| |
| // The nested list wraps as the level of operator -> subtask -> state object collection |
| Map<OperatorID, List<List<T>>> oldStates = |
| splitManagedAndRawOperatorStates(oldOperatorStates, extractHandle); |
| |
| oldOperatorStates.forEach( |
| (operatorID, oldOperatorState) -> |
| result.putAll( |
| applyRepartitioner( |
| operatorID, |
| stateRepartitioner, |
| oldStates.get(operatorID), |
| oldOperatorState.getParallelism(), |
| newParallelism))); |
| } |
| |
| public void reDistributeResultSubpartitionStates(TaskStateAssignment assignment) { |
| // FLINK-31963: We can skip this phase if there is no output state AND downstream has no |
| // input states |
| if (!assignment.hasOutputState && !assignment.hasDownstreamInputStates()) { |
| return; |
| } |
| |
| checkForUnsupportedToplogyChanges( |
| assignment.oldState, |
| OperatorSubtaskState::getResultSubpartitionState, |
| assignment.outputOperatorID); |
| |
| final OperatorState outputState = assignment.oldState.get(assignment.outputOperatorID); |
| final List<List<ResultSubpartitionStateHandle>> outputOperatorState = |
| splitBySubtasks(outputState, OperatorSubtaskState::getResultSubpartitionState); |
| |
| final ExecutionJobVertex executionJobVertex = assignment.executionJobVertex; |
| final List<IntermediateDataSet> outputs = |
| executionJobVertex.getJobVertex().getProducedDataSets(); |
| |
| if (outputState.getParallelism() == executionJobVertex.getParallelism()) { |
| assignment.resultSubpartitionStates.putAll( |
| toInstanceMap(assignment.outputOperatorID, outputOperatorState)); |
| return; |
| } |
| // Parallelism of this vertex changed, distribute ResultSubpartitionStateHandle |
| // according to output mapping. |
| for (int partitionIndex = 0; partitionIndex < outputs.size(); partitionIndex++) { |
| final List<List<ResultSubpartitionStateHandle>> partitionState = |
| outputs.size() == 1 |
| ? outputOperatorState |
| : getPartitionState( |
| outputOperatorState, |
| ResultSubpartitionInfo::getPartitionIdx, |
| partitionIndex); |
| final MappingBasedRepartitioner<ResultSubpartitionStateHandle> repartitioner = |
| new MappingBasedRepartitioner<>( |
| assignment.getOutputMapping(partitionIndex).getRescaleMappings()); |
| final Map<OperatorInstanceID, List<ResultSubpartitionStateHandle>> repartitioned = |
| applyRepartitioner( |
| assignment.outputOperatorID, |
| repartitioner, |
| partitionState, |
| outputOperatorState.size(), |
| executionJobVertex.getParallelism()); |
| addToSubtasks(assignment.resultSubpartitionStates, repartitioned); |
| } |
| } |
| |
| public void reDistributeInputChannelStates(TaskStateAssignment stateAssignment) { |
| // FLINK-31963: We can skip this phase only if there is no input state AND upstream has no |
| // output states |
| if (!stateAssignment.hasInputState && !stateAssignment.hasUpstreamOutputStates()) { |
| return; |
| } |
| |
| checkForUnsupportedToplogyChanges( |
| stateAssignment.oldState, |
| OperatorSubtaskState::getInputChannelState, |
| stateAssignment.inputOperatorID); |
| |
| final ExecutionJobVertex executionJobVertex = stateAssignment.executionJobVertex; |
| final List<IntermediateResult> inputs = executionJobVertex.getInputs(); |
| |
| // check for rescaling: no rescaling = simple reassignment |
| final OperatorState inputState = |
| stateAssignment.oldState.get(stateAssignment.inputOperatorID); |
| final List<List<InputChannelStateHandle>> inputOperatorState = |
| splitBySubtasks(inputState, OperatorSubtaskState::getInputChannelState); |
| |
| boolean hasAnyFullMapper = |
| executionJobVertex.getJobVertex().getInputs().stream() |
| .map(JobEdge::getDownstreamSubtaskStateMapper) |
| .anyMatch(m -> m.equals(SubtaskStateMapper.FULL)); |
| boolean hasAnyPreviousOperatorChanged = |
| executionJobVertex.getInputs().stream() |
| .map(IntermediateResult::getProducer) |
| .map(vertexAssignments::get) |
| .anyMatch( |
| taskStateAssignment -> { |
| final int oldParallelism = |
| stateAssignment |
| .oldState |
| .get(stateAssignment.inputOperatorID) |
| .getParallelism(); |
| return oldParallelism |
| != taskStateAssignment.executionJobVertex |
| .getParallelism(); |
| }); |
| |
| // need rescale if any input operator parallelism was changed and have any input with FULL |
| // subtask state mapper |
| if (inputState.getParallelism() == executionJobVertex.getParallelism() |
| && !(hasAnyFullMapper && hasAnyPreviousOperatorChanged)) { |
| stateAssignment.inputChannelStates.putAll( |
| toInstanceMap(stateAssignment.inputOperatorID, inputOperatorState)); |
| return; |
| } |
| // if this subtask has a different degree of parallelism, use the partitioner to figure |
| // out to which subtasks the state should be reassigned. In some cases, state is |
| // replicated to multiple subtasks and filtered during recovery. |
| // example: if this task is downscaled from 3 to 2 and it uses a range partitioner over |
| // [0;128) |
| // old assignment: 0 -> [0;43); 1 -> [43;87); 2 -> [87;128) |
| // new assignment: 0 -> [0;64]; 1 -> [64;128) |
| // subtask 0 recovers data from old subtask 0 + 1 and subtask 1 recovers data from old |
| // subtask 1 + 2 |
| for (int gateIndex = 0; gateIndex < inputs.size(); gateIndex++) { |
| final RescaleMappings mapping = |
| stateAssignment.getInputMapping(gateIndex).getRescaleMappings(); |
| |
| final List<List<InputChannelStateHandle>> gateState = |
| inputs.size() == 1 |
| ? inputOperatorState |
| : getPartitionState( |
| inputOperatorState, InputChannelInfo::getGateIdx, gateIndex); |
| final MappingBasedRepartitioner<InputChannelStateHandle> repartitioner = |
| new MappingBasedRepartitioner<>(mapping); |
| final Map<OperatorInstanceID, List<InputChannelStateHandle>> repartitioned = |
| applyRepartitioner( |
| stateAssignment.inputOperatorID, |
| repartitioner, |
| gateState, |
| inputOperatorState.size(), |
| stateAssignment.newParallelism); |
| addToSubtasks(stateAssignment.inputChannelStates, repartitioned); |
| } |
| } |
| |
| private static <K, V> void addToSubtasks(Map<K, List<V>> target, Map<K, List<V>> toAdd) { |
| toAdd.forEach( |
| (key, values) -> |
| target.computeIfAbsent(key, (unused) -> new ArrayList<>(values.size())) |
| .addAll(values)); |
| } |
| |
| private <T extends AbstractChannelStateHandle<?>> void checkForUnsupportedToplogyChanges( |
| Map<OperatorID, OperatorState> oldOperatorStates, |
| Function<OperatorSubtaskState, StateObjectCollection<T>> extractHandle, |
| OperatorID expectedOperatorID) { |
| final List<OperatorID> unexpectedState = |
| oldOperatorStates.entrySet().stream() |
| .filter(idAndState -> !idAndState.getKey().equals(expectedOperatorID)) |
| .filter(idAndState -> hasChannelState(idAndState.getValue(), extractHandle)) |
| .map(Map.Entry::getKey) |
| .collect(Collectors.toList()); |
| if (!unexpectedState.isEmpty()) { |
| throw new IllegalStateException( |
| "Cannot recover from unaligned checkpoint when topology changes, such that " |
| + "data exchanges with persisted data are now chained.\n" |
| + "The following operators contain channel state: " |
| + unexpectedState); |
| } |
| } |
| |
| private <T extends AbstractChannelStateHandle<?>> boolean hasChannelState( |
| OperatorState operatorState, |
| Function<OperatorSubtaskState, StateObjectCollection<T>> extractHandle) { |
| return operatorState.getSubtaskStates().values().stream() |
| .anyMatch(subState -> !isEmpty(extractHandle.apply(subState))); |
| } |
| |
| private <T extends AbstractChannelStateHandle<?>> boolean isEmpty(StateObjectCollection<T> s) { |
| return s.stream().allMatch(state -> state.getOffsets().isEmpty()); |
| } |
| |
| private static <T extends AbstractChannelStateHandle<I>, I> List<List<T>> getPartitionState( |
| List<List<T>> subtaskStates, Function<I, Integer> partitionExtractor, int partitionId) { |
| return subtaskStates.stream() |
| .map( |
| subtaskState -> |
| subtaskState.stream() |
| .filter( |
| state -> |
| partitionExtractor.apply(state.getInfo()) |
| == partitionId) |
| .collect(Collectors.toList())) |
| .collect(Collectors.toList()); |
| } |
| |
| private static <T extends StateObject> |
| Map<OperatorID, List<List<T>>> splitManagedAndRawOperatorStates( |
| Map<OperatorID, OperatorState> operatorStates, |
| Function<OperatorSubtaskState, StateObjectCollection<T>> extractHandle) { |
| return operatorStates.entrySet().stream() |
| .collect( |
| Collectors.toMap( |
| Map.Entry::getKey, |
| operatorIdAndState -> |
| splitBySubtasks( |
| operatorIdAndState.getValue(), extractHandle))); |
| } |
| |
| private static <T extends StateObject> List<List<T>> splitBySubtasks( |
| OperatorState operatorState, |
| Function<OperatorSubtaskState, StateObjectCollection<T>> extractHandle) { |
| List<List<T>> statePerSubtask = new ArrayList<>(operatorState.getParallelism()); |
| |
| for (int subTaskIndex = 0; subTaskIndex < operatorState.getParallelism(); subTaskIndex++) { |
| OperatorSubtaskState subtaskState = operatorState.getState(subTaskIndex); |
| statePerSubtask.add( |
| subtaskState == null |
| ? emptyList() |
| : extractHandle.apply(subtaskState).asList()); |
| } |
| |
| return statePerSubtask; |
| } |
| |
| /** |
| * Collect {@link KeyGroupsStateHandle managedKeyedStateHandles} which have intersection with |
| * given {@link KeyGroupRange} from {@link TaskState operatorState}. |
| * |
| * @param operatorState all state handles of a operator |
| * @param subtaskKeyGroupRange the KeyGroupRange of a subtask |
| * @return all managedKeyedStateHandles which have intersection with given KeyGroupRange |
| */ |
| public static List<KeyedStateHandle> getManagedKeyedStateHandles( |
| OperatorState operatorState, KeyGroupRange subtaskKeyGroupRange) { |
| |
| final int parallelism = operatorState.getParallelism(); |
| |
| List<KeyedStateHandle> subtaskKeyedStateHandles = null; |
| |
| for (int i = 0; i < parallelism; i++) { |
| if (operatorState.getState(i) != null) { |
| |
| Collection<KeyedStateHandle> keyedStateHandles = |
| operatorState.getState(i).getManagedKeyedState(); |
| |
| if (subtaskKeyedStateHandles == null) { |
| subtaskKeyedStateHandles = |
| new ArrayList<>(parallelism * keyedStateHandles.size()); |
| } |
| |
| extractIntersectingState( |
| keyedStateHandles, subtaskKeyGroupRange, subtaskKeyedStateHandles); |
| } |
| } |
| |
| return subtaskKeyedStateHandles != null ? subtaskKeyedStateHandles : emptyList(); |
| } |
| |
| /** |
| * Collect {@link KeyGroupsStateHandle rawKeyedStateHandles} which have intersection with given |
| * {@link KeyGroupRange} from {@link TaskState operatorState}. |
| * |
| * @param operatorState all state handles of a operator |
| * @param subtaskKeyGroupRange the KeyGroupRange of a subtask |
| * @return all rawKeyedStateHandles which have intersection with given KeyGroupRange |
| */ |
| public static List<KeyedStateHandle> getRawKeyedStateHandles( |
| OperatorState operatorState, KeyGroupRange subtaskKeyGroupRange) { |
| |
| final int parallelism = operatorState.getParallelism(); |
| |
| List<KeyedStateHandle> extractedKeyedStateHandles = null; |
| |
| for (int i = 0; i < parallelism; i++) { |
| if (operatorState.getState(i) != null) { |
| |
| Collection<KeyedStateHandle> rawKeyedState = |
| operatorState.getState(i).getRawKeyedState(); |
| |
| if (extractedKeyedStateHandles == null) { |
| extractedKeyedStateHandles = |
| new ArrayList<>(parallelism * rawKeyedState.size()); |
| } |
| |
| extractIntersectingState( |
| rawKeyedState, subtaskKeyGroupRange, extractedKeyedStateHandles); |
| } |
| } |
| |
| return extractedKeyedStateHandles != null ? extractedKeyedStateHandles : emptyList(); |
| } |
| |
| /** |
| * Extracts certain key group ranges from the given state handles and adds them to the |
| * collector. |
| */ |
| @VisibleForTesting |
| public static void extractIntersectingState( |
| Collection<? extends KeyedStateHandle> originalSubtaskStateHandles, |
| KeyGroupRange rangeToExtract, |
| List<KeyedStateHandle> extractedStateCollector) { |
| |
| for (KeyedStateHandle keyedStateHandle : originalSubtaskStateHandles) { |
| |
| if (keyedStateHandle != null) { |
| |
| KeyedStateHandle intersectedKeyedStateHandle = |
| keyedStateHandle.getIntersection(rangeToExtract); |
| |
| if (intersectedKeyedStateHandle != null) { |
| extractedStateCollector.add(intersectedKeyedStateHandle); |
| } |
| } |
| } |
| } |
| |
| /** |
| * Groups the available set of key groups into key group partitions. A key group partition is |
| * the set of key groups which is assigned to the same task. Each set of the returned list |
| * constitutes a key group partition. |
| * |
| * <p><b>IMPORTANT</b>: The assignment of key groups to partitions has to be in sync with the |
| * KeyGroupStreamPartitioner. |
| * |
| * @param numberKeyGroups Number of available key groups (indexed from 0 to numberKeyGroups - 1) |
| * @param parallelism Parallelism to generate the key group partitioning for |
| * @return List of key group partitions |
| */ |
| public static List<KeyGroupRange> createKeyGroupPartitions( |
| int numberKeyGroups, int parallelism) { |
| Preconditions.checkArgument(numberKeyGroups >= parallelism); |
| List<KeyGroupRange> result = new ArrayList<>(parallelism); |
| |
| for (int i = 0; i < parallelism; ++i) { |
| result.add( |
| KeyGroupRangeAssignment.computeKeyGroupRangeForOperatorIndex( |
| numberKeyGroups, parallelism, i)); |
| } |
| return result; |
| } |
| |
| /** |
| * Verifies conditions in regards to parallelism and maxParallelism that must be met when |
| * restoring state. |
| * |
| * @param operatorState state to restore |
| * @param executionJobVertex task for which the state should be restored |
| */ |
| private static void checkParallelismPreconditions( |
| OperatorState operatorState, ExecutionJobVertex executionJobVertex) { |
| // ----------------------------------------max parallelism |
| // preconditions------------------------------------- |
| |
| if (operatorState.getMaxParallelism() < executionJobVertex.getParallelism()) { |
| throw new IllegalStateException( |
| "The state for task " |
| + executionJobVertex.getJobVertexId() |
| + " can not be restored. The maximum parallelism (" |
| + operatorState.getMaxParallelism() |
| + ") of the restored state is lower than the configured parallelism (" |
| + executionJobVertex.getParallelism() |
| + "). Please reduce the parallelism of the task to be lower or equal to the maximum parallelism."); |
| } |
| |
| // check that the number of key groups have not changed or if we need to override it to |
| // satisfy the restored state |
| if (operatorState.getMaxParallelism() != executionJobVertex.getMaxParallelism()) { |
| |
| if (executionJobVertex.canRescaleMaxParallelism(operatorState.getMaxParallelism())) { |
| LOG.debug( |
| "Rescaling maximum parallelism for JobVertex {} from {} to {}", |
| executionJobVertex.getJobVertexId(), |
| executionJobVertex.getMaxParallelism(), |
| operatorState.getMaxParallelism()); |
| |
| executionJobVertex.setMaxParallelism(operatorState.getMaxParallelism()); |
| } else { |
| // if the max parallelism cannot be rescaled, we complain on mismatch |
| throw new IllegalStateException( |
| "The maximum parallelism (" |
| + operatorState.getMaxParallelism() |
| + ") with which the latest " |
| + "checkpoint of the execution job vertex " |
| + executionJobVertex |
| + " has been taken and the current maximum parallelism (" |
| + executionJobVertex.getMaxParallelism() |
| + ") changed. This " |
| + "is currently not supported."); |
| } |
| } |
| } |
| |
| /** |
| * Verifies that all operator states can be mapped to an execution job vertex. |
| * |
| * @param allowNonRestoredState if false an exception will be thrown if a state could not be |
| * mapped |
| * @param operatorStates operator states to map |
| * @param tasks task to map to |
| */ |
| private static void checkStateMappingCompleteness( |
| boolean allowNonRestoredState, |
| Map<OperatorID, OperatorState> operatorStates, |
| Set<ExecutionJobVertex> tasks) { |
| |
| Set<OperatorID> allOperatorIDs = new HashSet<>(); |
| for (ExecutionJobVertex executionJobVertex : tasks) { |
| for (OperatorIDPair operatorIDPair : executionJobVertex.getOperatorIDs()) { |
| allOperatorIDs.add(operatorIDPair.getGeneratedOperatorID()); |
| operatorIDPair.getUserDefinedOperatorID().ifPresent(allOperatorIDs::add); |
| } |
| } |
| for (Map.Entry<OperatorID, OperatorState> operatorGroupStateEntry : |
| operatorStates.entrySet()) { |
| // ----------------------------------------find operator for |
| // state--------------------------------------------- |
| |
| if (!allOperatorIDs.contains(operatorGroupStateEntry.getKey())) { |
| |
| OperatorState operatorState = operatorGroupStateEntry.getValue(); |
| |
| if (allowNonRestoredState) { |
| LOG.info( |
| "Skipped checkpoint state for operator {}.", |
| operatorState.getOperatorID()); |
| } else { |
| throw new IllegalStateException( |
| "There is no operator for the state " + operatorState.getOperatorID()); |
| } |
| } |
| } |
| } |
| |
| public static <T> Map<OperatorInstanceID, List<T>> applyRepartitioner( |
| OperatorID operatorID, |
| OperatorStateRepartitioner<T> opStateRepartitioner, |
| List<List<T>> chainOpParallelStates, |
| int oldParallelism, |
| int newParallelism) { |
| |
| List<List<T>> states = |
| applyRepartitioner( |
| opStateRepartitioner, |
| chainOpParallelStates, |
| oldParallelism, |
| newParallelism); |
| |
| return toInstanceMap(operatorID, states); |
| } |
| |
| private static <T> Map<OperatorInstanceID, List<T>> toInstanceMap( |
| OperatorID operatorID, List<List<T>> states) { |
| Map<OperatorInstanceID, List<T>> result = |
| CollectionUtil.newHashMapWithExpectedSize(states.size()); |
| |
| for (int subtaskIndex = 0; subtaskIndex < states.size(); subtaskIndex++) { |
| checkNotNull(states.get(subtaskIndex) != null, "states.get(subtaskIndex) is null"); |
| result.put(OperatorInstanceID.of(subtaskIndex, operatorID), states.get(subtaskIndex)); |
| } |
| |
| return result; |
| } |
| |
| /** |
| * Repartitions the given operator state using the given {@link OperatorStateRepartitioner} with |
| * respect to the new parallelism. |
| * |
| * @param opStateRepartitioner partitioner to use |
| * @param chainOpParallelStates state to repartition |
| * @param oldParallelism parallelism with which the state is currently partitioned |
| * @param newParallelism parallelism with which the state should be partitioned |
| * @return repartitioned state |
| */ |
| // TODO rewrite based on operator id |
| public static <T> List<List<T>> applyRepartitioner( |
| OperatorStateRepartitioner<T> opStateRepartitioner, |
| List<List<T>> chainOpParallelStates, |
| int oldParallelism, |
| int newParallelism) { |
| |
| if (chainOpParallelStates == null) { |
| return emptyList(); |
| } |
| |
| return opStateRepartitioner.repartitionState( |
| chainOpParallelStates, oldParallelism, newParallelism); |
| } |
| } |
