zhuzhurk commented on a change in pull request #12256:
URL: https://github.com/apache/flink/pull/12256#discussion_r429777663



##########
File path: 
flink-runtime/src/main/java/org/apache/flink/runtime/scheduler/DefaultExecutionSlotAllocator.java
##########
@@ -67,23 +76,211 @@
 
        private final SlotProviderStrategy slotProviderStrategy;
 
+       private final SlotOwner slotOwner;
+
        private final InputsLocationsRetriever inputsLocationsRetriever;
 
+       // temporary hack. can be removed along with the individual slot 
allocation code path
+       // once bulk slot allocation is fully functional
+       private final boolean enableBulkSlotAllocation;
+
        public DefaultExecutionSlotAllocator(
                        SlotProviderStrategy slotProviderStrategy,
-                       InputsLocationsRetriever inputsLocationsRetriever) {
+                       InputsLocationsRetriever inputsLocationsRetriever,
+                       boolean enableBulkSlotAllocation) {
                this.slotProviderStrategy = checkNotNull(slotProviderStrategy);
                this.inputsLocationsRetriever = 
checkNotNull(inputsLocationsRetriever);
+               this.enableBulkSlotAllocation = enableBulkSlotAllocation;
+
+               this.slotOwner = new DefaultExecutionSlotAllocatorSlotOwner();
 
                pendingSlotAssignments = new HashMap<>();
        }
 
+       /**
+        * Allocate slots in bulks. It has some limitations at the moment and 
is for pipelined region scheduling only.
+        * Slot sharing will be ignored.
+        * Co-location constraints are not allowed as a result.
+        * Intra-bulk input location preferences will be ignored.
+        */
+       private List<SlotExecutionVertexAssignment> allocateSlotsInBulks(
+                       final List<ExecutionVertexSchedulingRequirements> 
executionVertexSchedulingRequirements) {
+
+               
validateNoCoLocationConstraint(executionVertexSchedulingRequirements);
+
+               final Set<ExecutionVertexID> allExecutionVertexIds = 
executionVertexSchedulingRequirements.stream()
+                       
.map(ExecutionVertexSchedulingRequirements::getExecutionVertexId)
+                       .collect(Collectors.toSet());
+
+               final Map<ExecutionVertexID, SlotRequestId> 
executionVertexSlotRequestIds =
+                       
generateExecutionVertexSlotRequestIds(allExecutionVertexIds);
+
+               final List<SlotExecutionVertexAssignment> 
slotExecutionVertexAssignments =
+                       
createSlotExecutionVertexAssignments(executionVertexSchedulingRequirements, 
executionVertexSlotRequestIds);
+
+               final List<CompletableFuture<PhysicalSlotRequest>> 
physicalSlotRequestFutures =
+                       createPhysicalSlotRequestFutures(
+                               executionVertexSchedulingRequirements,
+                               allExecutionVertexIds,
+                               executionVertexSlotRequestIds);
+
+               allocateSlotsForAssignments(
+                       physicalSlotRequestFutures,
+                       slotExecutionVertexAssignments,
+                       executionVertexSlotRequestIds);
+
+               return slotExecutionVertexAssignments;
+       }
+
+       private static void validateNoCoLocationConstraint(
+                       final Collection<ExecutionVertexSchedulingRequirements> 
schedulingRequirements) {
+
+               final boolean hasCoLocationConstraint = 
schedulingRequirements.stream()
+                       .anyMatch(r -> r.getCoLocationConstraint() != null);
+               checkState(
+                       !hasCoLocationConstraint,
+                       "Jobs with co-location constraints are not allowed to 
run with pipelined region scheduling strategy.");
+       }
+
+       private Map<ExecutionVertexID, SlotRequestId> 
generateExecutionVertexSlotRequestIds(
+                       final Set<ExecutionVertexID> allExecutionVertexIds) {
+
+               final Map<ExecutionVertexID, SlotRequestId> 
executionVertexSlotRequestIds = new HashMap<>();
+               for (ExecutionVertexID executionVertexId : 
allExecutionVertexIds) {
+                       executionVertexSlotRequestIds.put(executionVertexId, 
new SlotRequestId());
+               }
+               return executionVertexSlotRequestIds;
+       }
+
+       private List<SlotExecutionVertexAssignment> 
createSlotExecutionVertexAssignments(
+                       final List<ExecutionVertexSchedulingRequirements> 
executionVertexSchedulingRequirements,
+                       final Map<ExecutionVertexID, SlotRequestId> 
executionVertexSlotRequestIds) {
+
+               final List<SlotExecutionVertexAssignment> 
slotExecutionVertexAssignments =
+                       new 
ArrayList<>(executionVertexSchedulingRequirements.size());
+               for (ExecutionVertexSchedulingRequirements 
schedulingRequirements : executionVertexSchedulingRequirements) {
+                       final ExecutionVertexID executionVertexId = 
schedulingRequirements.getExecutionVertexId();
+                       final SlotRequestId slotRequestId = 
executionVertexSlotRequestIds.get(executionVertexId);
+
+                       final SlotExecutionVertexAssignment 
slotExecutionVertexAssignment =
+                               createAndRegisterSlotExecutionVertexAssignment(
+                                       executionVertexId,
+                                       new CompletableFuture<>(),
+                                       slotRequestId,
+                                       null);
+                       
slotExecutionVertexAssignments.add(slotExecutionVertexAssignment);
+               }
+
+               return slotExecutionVertexAssignments;
+       }
+
+       private List<CompletableFuture<PhysicalSlotRequest>> 
createPhysicalSlotRequestFutures(
+                       final List<ExecutionVertexSchedulingRequirements> 
executionVertexSchedulingRequirements,
+                       final Set<ExecutionVertexID> allExecutionVertexIds,
+                       final Map<ExecutionVertexID, SlotRequestId> 
executionVertexSlotRequestIds) {
+
+               final Set<AllocationID> allPreviousAllocationIds =
+                       
computeAllPriorAllocationIds(executionVertexSchedulingRequirements);
+
+               final List<CompletableFuture<PhysicalSlotRequest>> 
physicalSlotRequestFutures =
+                       new 
ArrayList<>(executionVertexSchedulingRequirements.size());
+               for (ExecutionVertexSchedulingRequirements 
schedulingRequirements : executionVertexSchedulingRequirements) {
+                       final ExecutionVertexID executionVertexId = 
schedulingRequirements.getExecutionVertexId();
+                       final SlotRequestId slotRequestId = 
executionVertexSlotRequestIds.get(executionVertexId);
+
+                       LOG.debug("Allocate slot with id {} for execution {}", 
slotRequestId, executionVertexId);
+
+                       // use the task resource profile as the physical slot 
resource requirement since slot sharing is ignored here
+                       final ResourceProfile physicalSlotResourceProfile = 
schedulingRequirements.getTaskResourceProfile();
+
+                       final CompletableFuture<SlotProfile> slotProfileFuture 
= createSlotProfile(
+                               schedulingRequirements,
+                               allExecutionVertexIds,
+                               physicalSlotResourceProfile,
+                               allPreviousAllocationIds);
+                       final CompletableFuture<PhysicalSlotRequest> 
physicalSlotRequestFuture =
+                               slotProfileFuture.thenApply(
+                                       slotProfile -> 
createPhysicalSlotRequest(slotRequestId, slotProfile));
+                       
physicalSlotRequestFutures.add(physicalSlotRequestFuture);
+               }
+
+               return physicalSlotRequestFutures;
+       }
+
+       private PhysicalSlotRequest createPhysicalSlotRequest(
+                       final SlotRequestId slotRequestId,
+                       final SlotProfile slotProfile) {
+               return new PhysicalSlotRequest(slotRequestId, slotProfile, 
slotProviderStrategy.willSlotBeOccupiedIndefinitely());
+       }
+
+       private void allocateSlotsForAssignments(
+                       final List<CompletableFuture<PhysicalSlotRequest>> 
physicalSlotRequestFutures,
+                       final List<SlotExecutionVertexAssignment> 
slotExecutionVertexAssignments,
+                       final Map<ExecutionVertexID, SlotRequestId> 
executionVertexSlotRequestIds) {
+
+               final Map<ExecutionVertexID, SlotExecutionVertexAssignment> 
vertexToAssignmentMapping =
+                       slotExecutionVertexAssignments.stream()
+                               
.collect(Collectors.toMap(SlotExecutionVertexAssignment::getExecutionVertexId, 
Function.identity()));
+
+               final Map<SlotRequestId, ExecutionVertexID> slotToVertexMapping 
= new HashMap<>();
+               for (Map.Entry<ExecutionVertexID, SlotRequestId> entry : 
executionVertexSlotRequestIds.entrySet()) {
+                       slotToVertexMapping.put(entry.getValue(), 
entry.getKey());
+               }
+
+               FutureUtils.combineAll(physicalSlotRequestFutures)
+                       .thenCompose(physicalSlotRequests -> 
slotProviderStrategy.allocatePhysicalSlots(physicalSlotRequests))
+                       .thenAccept(physicalSlotRequestResults -> {
+                               for (PhysicalSlotRequest.Result result : 
physicalSlotRequestResults) {
+                                       final SlotRequestId slotRequestId = 
result.getSlotRequestId();
+                                       final ExecutionVertexID 
executionVertexId = slotToVertexMapping.get(slotRequestId);
+                                       try {
+                                               final LogicalSlot logicalSlot = 
allocateSingleLogicalSlotFromPhysicalSlot(
+                                                       slotRequestId,
+                                                       
result.getPhysicalSlot());
+                                               
vertexToAssignmentMapping.get(executionVertexId).getLogicalSlotFuture().complete(logicalSlot);
+                                       } catch (Exception ex) {
+                                               throw new 
CompletionException(ex);
+                                       }
+                               }
+                       })
+                       .exceptionally(ex -> {
+                               executionVertexSlotRequestIds.values().stream()
+                                       .forEach(slotRequestId -> {
+                                               final ExecutionVertexID 
executionVertexId = slotToVertexMapping.get(slotRequestId);
+                                               
vertexToAssignmentMapping.get(executionVertexId).getLogicalSlotFuture().completeExceptionally(ex);
+                                       });
+                               return null;
+                       });
+       }
+
+       private LogicalSlot allocateSingleLogicalSlotFromPhysicalSlot(
+                       final SlotRequestId slotRequestId,
+                       final PhysicalSlot physicalSlot) throws FlinkException {
+
+               final SingleLogicalSlot singleTaskSlot = new SingleLogicalSlot(
+                       slotRequestId,
+                       physicalSlot,
+                       null,
+                       Locality.UNKNOWN,
+                       slotOwner);
+
+               if (physicalSlot.tryAssignPayload(singleTaskSlot)) {
+                       return singleTaskSlot;
+               } else {
+                       throw new FlinkException("Could not assign payload to 
allocated slot " + physicalSlot.getAllocationId() + '.');
+               }
+       }
+
        @Override
        public List<SlotExecutionVertexAssignment> allocateSlotsFor(

Review comment:
       Done with methods of `OneSlotPerExecutionSlotAllocator` moved out.




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
[email protected]


Reply via email to