Examples with RemoteTask io.prestosql.execution.RemoteTask used on opensource projects

Example 16 with RemoteTask

use of io.prestosql.execution.RemoteTask in project hetu-core by openlookeng.

the class SourcePartitionedScheduler method schedule.

@Override
public synchronized ScheduleResult schedule(int maxSplitGroup) {
    dropListenersFromWhenFinishedOrNewLifespansAdded();
    int overallSplitAssignmentCount = 0;
    ImmutableSet.Builder<RemoteTask> overallNewTasks = ImmutableSet.builder();
    List<ListenableFuture<?>> overallBlockedFutures = new ArrayList<>();
    boolean anyBlockedOnPlacements = false;
    boolean anyBlockedOnNextSplitBatch = false;
    boolean anyNotBlocked = false;
    boolean applyFilter = isHeuristicIndexFilterEnabled(session) && SplitFiltering.isSplitFilterApplicable(stage);
    boolean initialMarker = false;
    for (Entry<Lifespan, ScheduleGroup> entry : scheduleGroups.entrySet()) {
        Lifespan lifespan = entry.getKey();
        ScheduleGroup scheduleGroup = entry.getValue();
        Set<Split> pendingSplits = scheduleGroup.pendingSplits;
        if (scheduleGroup.state == ScheduleGroupState.NO_MORE_SPLITS || scheduleGroup.state == ScheduleGroupState.DONE) {
            verify(scheduleGroup.nextSplitBatchFuture == null);
        } else if (pendingSplits.isEmpty()) {
            // try to get the next batch
            if (scheduleGroup.nextSplitBatchFuture == null) {
                scheduleGroup.nextSplitBatchFuture = splitSource.getNextBatch(scheduleGroup.partitionHandle, lifespan, splitBatchSize - pendingSplits.size());
                long start = System.nanoTime();
                addSuccessCallback(scheduleGroup.nextSplitBatchFuture, () -> stage.recordGetSplitTime(start));
            }
            if (scheduleGroup.nextSplitBatchFuture.isDone()) {
                SplitBatch nextSplits = getFutureValue(scheduleGroup.nextSplitBatchFuture);
                scheduleGroup.nextSplitBatchFuture = null;
                // add split filter to filter out split has no valid rows
                Pair<Optional<RowExpression>, Map<Symbol, ColumnHandle>> pair = SplitFiltering.getExpression(stage);
                if (SystemSessionProperties.isSnapshotEnabled(session)) {
                    List<Split> batchSplits = nextSplits.getSplits();
                    // Don't apply filter to MarkerSplit
                    if (batchSplits.size() == 1 && batchSplits.get(0).getConnectorSplit() instanceof MarkerSplit) {
                        applyFilter = false;
                    }
                }
                List<Split> filteredSplit = applyFilter ? SplitFiltering.getFilteredSplit(pair.getFirst(), SplitFiltering.getFullyQualifiedName(stage), pair.getSecond(), nextSplits, heuristicIndexerManager) : nextSplits.getSplits();
                // In case of ORC small size files/splits are grouped
                List<Split> groupedSmallFilesList = splitSource.groupSmallSplits(filteredSplit, lifespan, maxSplitGroup);
                filteredSplit = groupedSmallFilesList;
                pendingSplits.addAll(filteredSplit);
                if (nextSplits.isLastBatch()) {
                    if (scheduleGroup.state == ScheduleGroupState.INITIALIZED && pendingSplits.isEmpty()) {
                        // Add an empty split in case no splits have been produced for the source.
                        // For source operators, they never take input, but they may produce output.
                        // This is well handled by Presto execution engine.
                        // However, there are certain non-source operators that may produce output without any input,
                        // for example, 1) an AggregationOperator, 2) a HashAggregationOperator where one of the grouping sets is ().
                        // Scheduling an empty split kicks off necessary driver instantiation to make this work.
                        pendingSplits.add(new Split(splitSource.getCatalogName(), new EmptySplit(splitSource.getCatalogName()), lifespan));
                    }
                    scheduleGroup.state = ScheduleGroupState.NO_MORE_SPLITS;
                }
            } else {
                overallBlockedFutures.add(scheduleGroup.nextSplitBatchFuture);
                anyBlockedOnNextSplitBatch = true;
                continue;
            }
        }
        Multimap<InternalNode, Split> splitAssignment = ImmutableMultimap.of();
        if (!pendingSplits.isEmpty()) {
            if (!scheduleGroup.placementFuture.isDone()) {
                anyBlockedOnPlacements = true;
                continue;
            }
            if (scheduleGroup.state == ScheduleGroupState.INITIALIZED) {
                scheduleGroup.state = ScheduleGroupState.SPLITS_ADDED;
            }
            if (state == State.INITIALIZED) {
                state = State.SPLITS_ADDED;
            }
            // calculate placements for splits
            SplitPlacementResult splitPlacementResult;
            if (stage.isThrottledSchedule()) {
                // If asked for partial schedule incase of lesser resource, then schedule only 10% of splits.
                // 10% is calculated on initial number of splits and same is being used on subsequent schedule also.
                // But if later 10% of current pending splits more than earlier 10%, then it will schedule max of
                // these.
                // if throttledSplitsCount is more than number of pendingSplits, then it will schedule all.
                throttledSplitsCount = Math.max((int) Math.ceil(pendingSplits.size() * ALLOWED_PERCENT_LIMIT), throttledSplitsCount);
                splitPlacementResult = splitPlacementPolicy.computeAssignments(ImmutableSet.copyOf(Iterables.limit(pendingSplits, throttledSplitsCount)), this.stage);
            } else {
                splitPlacementResult = splitPlacementPolicy.computeAssignments(new HashSet<>(pendingSplits), this.stage);
            }
            splitAssignment = splitPlacementResult.getAssignments();
            if (SystemSessionProperties.isSnapshotEnabled(session)) {
                Split firstSplit = pendingSplits.iterator().next();
                if (pendingSplits.size() == 1 && firstSplit.getConnectorSplit() instanceof MarkerSplit) {
                    // We'll create a new assignment, but still need to call computeAssignments above, and cannot modify the returned assignment map directly
                    splitAssignment = HashMultimap.create(splitAssignment);
                    splitAssignment.values().remove(firstSplit);
                    // Getting all internalNodes and assigning marker splits to all of them.
                    List<InternalNode> allNodes = splitPlacementPolicy.allNodes();
                    for (InternalNode node : allNodes) {
                        splitAssignment.put(node, firstSplit);
                    }
                    MarkerSplit markerSplit = (MarkerSplit) firstSplit.getConnectorSplit();
                    // then set the flag below to true, so stages enter SCHEDULING_SPLITS state.
                    if (markerSplit.isResuming() || markerSplit.getSnapshotId() == 0) {
                        initialMarker = true;
                    }
                } else {
                    // MarkerSplit should be in its own batch.
                    verify(pendingSplits.stream().noneMatch(split -> split.getConnectorSplit() instanceof MarkerSplit));
                }
            }
            // remove splits with successful placements
            // AbstractSet.removeAll performs terribly here.
            splitAssignment.values().forEach(pendingSplits::remove);
            overallSplitAssignmentCount += splitAssignment.size();
            // if not completed placed, mark scheduleGroup as blocked on placement
            if (!pendingSplits.isEmpty()) {
                scheduleGroup.placementFuture = splitPlacementResult.getBlocked();
                overallBlockedFutures.add(scheduleGroup.placementFuture);
                anyBlockedOnPlacements = true;
            }
        }
        // if no new splits will be assigned, update state and attach completion event
        Multimap<InternalNode, Lifespan> noMoreSplitsNotification = ImmutableMultimap.of();
        if (pendingSplits.isEmpty() && scheduleGroup.state == ScheduleGroupState.NO_MORE_SPLITS) {
            scheduleGroup.state = ScheduleGroupState.DONE;
            if (!lifespan.isTaskWide()) {
                InternalNode node = ((BucketedSplitPlacementPolicy) splitPlacementPolicy).getNodeForBucket(lifespan.getId());
                noMoreSplitsNotification = ImmutableMultimap.of(node, lifespan);
            }
        }
        // assign the splits with successful placements
        overallNewTasks.addAll(assignSplits(splitAssignment, noMoreSplitsNotification));
        // As a result, to avoid busy loops caused by 1, we check pendingSplits.isEmpty() instead of placementFuture.isDone() here.
        if (scheduleGroup.nextSplitBatchFuture == null && scheduleGroup.pendingSplits.isEmpty() && scheduleGroup.state != ScheduleGroupState.DONE) {
            anyNotBlocked = true;
        }
    }
    // Next time it invokes getNextBatch, it will realize that. However, the invocation will fail we tear down splitSource now.
    if ((state == State.NO_MORE_SPLITS || state == State.FINISHED) || (noMoreScheduleGroups && scheduleGroups.isEmpty() && splitSource.isFinished())) {
        switch(state) {
            case INITIALIZED:
                // But this shouldn't be possible. See usage of EmptySplit in this method.
                throw new IllegalStateException("At least 1 split should have been scheduled for this plan node");
            case SPLITS_ADDED:
                state = State.NO_MORE_SPLITS;
                splitSource.close();
            // fall through
            case NO_MORE_SPLITS:
                state = State.FINISHED;
                whenFinishedOrNewLifespanAdded.set(null);
            // fall through
            case FINISHED:
                return new ScheduleResult(true, overallNewTasks.build(), overallSplitAssignmentCount);
            default:
                throw new IllegalStateException("Unknown state");
        }
    }
    if (anyNotBlocked) {
        if (initialMarker) {
            stage.transitionToSchedulingSplits();
        }
        return new ScheduleResult(false, overallNewTasks.build(), overallSplitAssignmentCount);
    }
    if (anyBlockedOnPlacements || groupedExecution) {
        // In a broadcast join, output buffers of the tasks in build source stage have to
        // hold onto all data produced before probe side task scheduling finishes,
        // even if the data is acknowledged by all known consumers. This is because
        // new consumers may be added until the probe side task scheduling finishes.
        // 
        // As a result, the following line is necessary to prevent deadlock
        // due to neither build nor probe can make any progress.
        // The build side blocks due to a full output buffer.
        // In the meantime the probe side split cannot be consumed since
        // builder side hash table construction has not finished.
        overallNewTasks.addAll(finalizeTaskCreationIfNecessary());
    }
    ScheduleResult.BlockedReason blockedReason;
    if (anyBlockedOnNextSplitBatch) {
        blockedReason = anyBlockedOnPlacements ? MIXED_SPLIT_QUEUES_FULL_AND_WAITING_FOR_SOURCE : WAITING_FOR_SOURCE;
    } else {
        blockedReason = anyBlockedOnPlacements ? SPLIT_QUEUES_FULL : NO_ACTIVE_DRIVER_GROUP;
    }
    overallBlockedFutures.add(whenFinishedOrNewLifespanAdded);
    return new ScheduleResult(false, overallNewTasks.build(), nonCancellationPropagating(whenAnyComplete(overallBlockedFutures)), blockedReason, overallSplitAssignmentCount);
}

Example 17 with RemoteTask

use of io.prestosql.execution.RemoteTask in project hetu-core by openlookeng.

the class ScaledWriterScheduler method scheduleTasks.

private List<RemoteTask> scheduleTasks(int count) {
    if (count == 0) {
        return ImmutableList.of();
    }
    List<InternalNode> nodes = nodeSelector.selectRandomNodes(count, scheduledNodes);
    if (scheduledNodes.isEmpty() && nodes.size() != initialTaskCount) {
        checkCondition(false, NO_NODES_AVAILABLE, "No nodes available to run query");
    }
    ImmutableList.Builder<RemoteTask> tasks = ImmutableList.builder();
    for (InternalNode node : nodes) {
        Optional<RemoteTask> remoteTask = stage.scheduleTask(node, scheduledNodes.size(), OptionalInt.empty());
        remoteTask.ifPresent(task -> {
            tasks.add(task);
            scheduledNodes.add(node);
        });
    }
    return tasks.build();
}

Example 18 with RemoteTask

use of io.prestosql.execution.RemoteTask in project hetu-core by openlookeng.

the class NodeScheduler method toWhenHasSplitQueueSpaceFuture.

public static ListenableFuture<?> toWhenHasSplitQueueSpaceFuture(Set<InternalNode> blockedNodes, List<RemoteTask> existingTasks, int spaceThreshold) {
    if (blockedNodes.isEmpty()) {
        return immediateFuture(null);
    }
    Map<String, RemoteTask> nodeToTaskMap = new HashMap<>();
    for (RemoteTask task : existingTasks) {
        nodeToTaskMap.put(task.getNodeId(), task);
    }
    List<ListenableFuture<?>> blockedFutures = blockedNodes.stream().map(InternalNode::getNodeIdentifier).map(nodeToTaskMap::get).filter(Objects::nonNull).map(remoteTask -> remoteTask.whenSplitQueueHasSpace(spaceThreshold)).collect(toImmutableList());
    if (blockedFutures.isEmpty()) {
        return immediateFuture(null);
    }
    return whenAnyCompleteCancelOthers(blockedFutures);
}

Example 19 with RemoteTask

use of io.prestosql.execution.RemoteTask in project hetu-core by openlookeng.

the class FixedCountScheduler method schedule.

@Override
public ScheduleResult schedule() {
    OptionalInt totalPartitions = OptionalInt.of(partitionToNode.size());
    List<RemoteTask> newTasks = IntStream.range(0, partitionToNode.size()).mapToObj(partition -> taskScheduler.scheduleTask(partitionToNode.get(partition), partition, totalPartitions)).filter(Optional::isPresent).map(Optional::get).collect(toImmutableList());
    return new ScheduleResult(true, newTasks, 0);
}

Example 20 with RemoteTask

use of io.prestosql.execution.RemoteTask in project hetu-core by openlookeng.

the class TestNodeScheduler method testSplitCount.

@Test
public void testSplitCount() {
    setUpNodes();
    MockRemoteTaskFactory remoteTaskFactory = new MockRemoteTaskFactory(remoteTaskExecutor, remoteTaskScheduledExecutor);
    InternalNode chosenNode = Iterables.get(nodeManager.getActiveConnectorNodes(CONNECTOR_ID), 0);
    TaskId taskId1 = new TaskId("test", 1, 1);
    RemoteTask remoteTask1 = remoteTaskFactory.createTableScanTask(taskId1, chosenNode, ImmutableList.of(new Split(CONNECTOR_ID, new TestSplitRemote(), Lifespan.taskWide()), new Split(CONNECTOR_ID, new TestSplitRemote(), Lifespan.taskWide())), nodeTaskMap.createPartitionedSplitCountTracker(chosenNode, taskId1));
    TaskId taskId2 = new TaskId("test", 1, 2);
    RemoteTask remoteTask2 = remoteTaskFactory.createTableScanTask(taskId2, chosenNode, ImmutableList.of(new Split(CONNECTOR_ID, new TestSplitRemote(), Lifespan.taskWide())), nodeTaskMap.createPartitionedSplitCountTracker(chosenNode, taskId2));
    nodeTaskMap.addTask(chosenNode, remoteTask1);
    nodeTaskMap.addTask(chosenNode, remoteTask2);
    assertEquals(nodeTaskMap.getPartitionedSplitsOnNode(chosenNode), 3);
    remoteTask1.abort();
    assertEquals(nodeTaskMap.getPartitionedSplitsOnNode(chosenNode), 1);
    remoteTask2.abort();
    assertEquals(nodeTaskMap.getPartitionedSplitsOnNode(chosenNode), 0);
}