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

Example 11 with RemoteTask

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

the class SplitCacheAwareNodeSelector method computeAssignments.

@Override
public SplitPlacementResult computeAssignments(Set<Split> splits, List<RemoteTask> existingTasks, Optional<SqlStageExecution> stage) {
    Multimap<InternalNode, Split> assignment = HashMultimap.create();
    NodeMap nodeMapSlice = this.nodeMap.get().get();
    Map<CatalogName, Map<String, InternalNode>> activeNodesByCatalog = new HashMap<>();
    NodeAssignmentStats assignmentStats = new NodeAssignmentStats(nodeTaskMap, nodeMapSlice, existingTasks);
    Set<Split> uncacheableSplits = new HashSet<>();
    Set<Split> newCacheableSplits = new HashSet<>();
    SplitCacheMap splitCacheMap = SplitCacheMap.getInstance();
    for (Split split : splits) {
        Optional<String> assignedNodeId = Optional.empty();
        SplitKey splitKey = createSplitKey(split);
        if (splitKey != null) {
            assignedNodeId = splitCacheMap.getCachedNodeId(splitKey);
        }
        if (!split.getConnectorSplit().isCacheable() || splitKey == null) {
            // uncacheable splits will be scheduled using default node selector
            uncacheableSplits.add(split);
            continue;
        }
        Map<String, InternalNode> activeNodes = activeNodesByCatalog.computeIfAbsent(split.getCatalogName(), catalogName -> nodeManager.getActiveConnectorNodes(catalogName).stream().collect(Collectors.toMap(InternalNode::getNodeIdentifier, Function.identity())));
        InternalNode assignedNode = assignedNodeId.map(activeNodes::get).orElse(null);
        // check if a node has been assigned and ensure it is still active before scheduling
        if (assignedNode != null) {
            // split has been previously assigned to a node
            // assign the split to the same node as before
            assignment.put(assignedNode, split);
            assignmentStats.addAssignedSplit(assignedNode);
        } else {
            // splits that have not be previously cached or the assigned node is now inactive
            newCacheableSplits.add(split);
        }
    }
    log.info("%d out of %d splits already cached. %d new splits to be cached. %d splits cannot be cached.", assignment.size(), splits.size(), newCacheableSplits.size(), uncacheableSplits.size());
    Set<Split> unassignedSplits = new HashSet<>();
    unassignedSplits.addAll(newCacheableSplits);
    unassignedSplits.addAll(uncacheableSplits);
    // Compute split assignments for splits that cannot be cached, newly cacheable, and already cached but cached worker is inactive now.
    SplitPlacementResult defaultSplitPlacementResult = defaultNodeSelector.computeAssignments(unassignedSplits, existingTasks, stage);
    defaultSplitPlacementResult.getAssignments().forEach(((internalNode, split) -> {
        // Set or Update cached node id only if split is cacheable
        if (newCacheableSplits.contains(split)) {
            SplitKey splitKey = createSplitKey(split);
            if (splitKey != null) {
                splitCacheMap.addCachedNode(splitKey, internalNode.getNodeIdentifier());
            }
        }
        assignmentStats.addAssignedSplit(internalNode);
    }));
    assignment.putAll(defaultSplitPlacementResult.getAssignments());
    // Check if its CTE node and its feeder
    if (stage.isPresent() && stage.get().getFragment().getFeederCTEId().isPresent()) {
        updateFeederNodeAndSplitCount(stage.get(), assignment);
    }
    return new SplitPlacementResult(defaultSplitPlacementResult.getBlocked(), assignment);
}

Example 12 with RemoteTask

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

the class SqlQueryScheduler method updateQueryOutputLocations.

private static void updateQueryOutputLocations(QueryStateMachine queryStateMachine, OutputBufferId rootBufferId, Set<RemoteTask> tasks, boolean noMoreExchangeLocations) {
    Set<TaskLocation> bufferLocations = tasks.stream().map(task -> {
        URI uri = task.getTaskStatus().getSelf();
        String instanceId = task.getInstanceId();
        URI newUri = uriBuilderFrom(uri).appendPath("results").appendPath(rootBufferId.toString()).build();
        return new TaskLocation(newUri, instanceId);
    }).collect(toImmutableSet());
    queryStateMachine.updateOutputLocations(bufferLocations, noMoreExchangeLocations);
}

Example 13 with RemoteTask

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

the class FixedSourcePartitionedScheduler method schedule.

@Override
public ScheduleResult schedule(int maxSplitGroup) {
    // schedule a task on every node in the distribution
    List<RemoteTask> newTasks = ImmutableList.of();
    if (!scheduledTasks) {
        OptionalInt totalPartitions = OptionalInt.of(nodes.size());
        newTasks = Streams.mapWithIndex(nodes.stream(), (node, id) -> stage.scheduleTask(node, toIntExact(id), totalPartitions)).filter(Optional::isPresent).map(Optional::get).collect(toImmutableList());
        scheduledTasks = true;
    }
    boolean allBlocked = true;
    List<ListenableFuture<?>> blocked = new ArrayList<>();
    BlockedReason blockedReason = BlockedReason.NO_ACTIVE_DRIVER_GROUP;
    if (groupedLifespanScheduler.isPresent()) {
        // Start new driver groups on the first scheduler if necessary,
        // i.e. when previous ones have finished execution (not finished scheduling).
        // 
        // Invoke schedule method to get a new SettableFuture every time.
        // Reusing previously returned SettableFuture could lead to the ListenableFuture retaining too many listeners.
        blocked.add(groupedLifespanScheduler.get().schedule(sourceSchedulers.get(0)));
    }
    int splitsScheduled = 0;
    Iterator<SourceScheduler> schedulerIterator = sourceSchedulers.iterator();
    List<Lifespan> driverGroupsToStart = ImmutableList.of();
    boolean shouldInvokeNoMoreDriverGroups = false;
    while (schedulerIterator.hasNext()) {
        SourceScheduler sourceScheduler = schedulerIterator.next();
        for (Lifespan lifespan : driverGroupsToStart) {
            sourceScheduler.startLifespan(lifespan, partitionHandleFor(lifespan));
        }
        if (shouldInvokeNoMoreDriverGroups) {
            sourceScheduler.noMoreLifespans();
        }
        ScheduleResult schedule = sourceScheduler.schedule(maxSplitGroup);
        splitsScheduled += schedule.getSplitsScheduled();
        if (schedule.getBlockedReason().isPresent()) {
            blocked.add(schedule.getBlocked());
            blockedReason = blockedReason.combineWith(schedule.getBlockedReason().get());
        } else {
            verify(schedule.getBlocked().isDone(), "blockedReason not provided when scheduler is blocked");
            allBlocked = false;
        }
        driverGroupsToStart = sourceScheduler.drainCompletedLifespans();
        if (schedule.isFinished()) {
            stage.schedulingComplete(sourceScheduler.getPlanNodeId());
            schedulerIterator.remove();
            sourceScheduler.close();
            shouldInvokeNoMoreDriverGroups = true;
        } else {
            shouldInvokeNoMoreDriverGroups = false;
        }
    }
    if (allBlocked) {
        return new ScheduleResult(sourceSchedulers.isEmpty(), newTasks, whenAnyComplete(blocked), blockedReason, splitsScheduled);
    } else {
        return new ScheduleResult(sourceSchedulers.isEmpty(), newTasks, splitsScheduled);
    }
}

Example 14 with RemoteTask

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

the class SqlQueryScheduler method createStages.

private List<SqlStageExecution> createStages(ExchangeLocationsConsumer parent, AtomicInteger nextStageId, LocationFactory locationFactory, StageExecutionPlan plan, NodeScheduler nodeScheduler, RemoteTaskFactory remoteTaskFactory, Session session, int splitBatchSize, BiFunction<PartitioningHandle, Integer, NodePartitionMap> partitioningCache, NodePartitioningManager nodePartitioningManager, ExecutorService queryExecutor, ScheduledExecutorService schedulerExecutor, FailureDetector failureDetector, NodeTaskMap nodeTaskMap, ImmutableMap.Builder<StageId, StageScheduler> stageSchedulers, ImmutableMap.Builder<StageId, StageLinkage> stageLinkages, boolean isSnapshotEnabled, QuerySnapshotManager snapshotManager, Map<StageId, Integer> stageTaskCounts) {
    ImmutableList.Builder<SqlStageExecution> localStages = ImmutableList.builder();
    StageId stageId = new StageId(queryStateMachine.getQueryId(), nextStageId.getAndIncrement());
    SqlStageExecution stageExecution = createSqlStageExecution(stageId, locationFactory.createStageLocation(stageId), plan.getFragment(), plan.getTables(), remoteTaskFactory, session, summarizeTaskInfo, nodeTaskMap, queryExecutor, failureDetector, schedulerStats, dynamicFilterService, snapshotManager);
    localStages.add(stageExecution);
    Optional<int[]> bucketToPartition;
    PartitioningHandle partitioningHandle = plan.getFragment().getPartitioning();
    boolean keepConsumerOnFeederNodes = !plan.getFragment().getFeederCTEId().isPresent() && plan.getFragment().getFeederCTEParentId().isPresent();
    if (partitioningHandle.equals(SOURCE_DISTRIBUTION)) {
        // nodes are selected dynamically based on the constraints of the splits and the system load
        Entry<PlanNodeId, SplitSource> entry = Iterables.getOnlyElement(plan.getSplitSources().entrySet());
        PlanNodeId planNodeId = entry.getKey();
        SplitSource splitSource = entry.getValue();
        CatalogName catalogName = splitSource.getCatalogName();
        if (isInternalSystemConnector(catalogName)) {
            catalogName = null;
        }
        NodeSelector nodeSelector = nodeScheduler.createNodeSelector(catalogName, keepConsumerOnFeederNodes, feederScheduledNodes);
        if (isSnapshotEnabled) {
            // When snapshot is enabled, then no task can be added after the query started running,
            // otherwise assumptions about how many "input channels" may be broken.
            nodeSelector.lockDownNodes();
        }
        SplitPlacementPolicy placementPolicy = new DynamicSplitPlacementPolicy(nodeSelector, stageExecution::getAllTasks);
        checkArgument(!plan.getFragment().getStageExecutionDescriptor().isStageGroupedExecution());
        stageSchedulers.put(stageId, newSourcePartitionedSchedulerAsStageScheduler(stageExecution, planNodeId, splitSource, placementPolicy, splitBatchSize, session, heuristicIndexerManager));
        bucketToPartition = Optional.of(new int[1]);
    } else if (partitioningHandle.equals(SCALED_WRITER_DISTRIBUTION)) {
        bucketToPartition = Optional.of(new int[1]);
    } else {
        Map<PlanNodeId, SplitSource> splitSources = plan.getSplitSources();
        if (!splitSources.isEmpty()) {
            // contains local source
            List<PlanNodeId> schedulingOrder = plan.getFragment().getPartitionedSources();
            CatalogName catalogName = partitioningHandle.getConnectorId().orElseThrow(IllegalStateException::new);
            List<ConnectorPartitionHandle> connectorPartitionHandles;
            boolean groupedExecutionForStage = plan.getFragment().getStageExecutionDescriptor().isStageGroupedExecution();
            if (groupedExecutionForStage) {
                connectorPartitionHandles = nodePartitioningManager.listPartitionHandles(session, partitioningHandle);
                checkState(!ImmutableList.of(NOT_PARTITIONED).equals(connectorPartitionHandles));
            } else {
                connectorPartitionHandles = ImmutableList.of(NOT_PARTITIONED);
            }
            BucketNodeMap bucketNodeMap;
            List<InternalNode> stageNodeList;
            if (plan.getFragment().getRemoteSourceNodes().stream().allMatch(node -> node.getExchangeType() == REPLICATE)) {
                // no remote source
                boolean dynamicLifespanSchedule = plan.getFragment().getStageExecutionDescriptor().isDynamicLifespanSchedule();
                if (isSnapshotEnabled) {
                    NodeSelector nodeSelector = nodeScheduler.createNodeSelector(catalogName, keepConsumerOnFeederNodes, feederScheduledNodes);
                    int nodeCount;
                    if (stageTaskCounts != null) {
                        // Resuming: need to create same number of tasks as old stage.
                        nodeCount = stageTaskCounts.get(stageId);
                    } else {
                        // Scheduling: reserve some nodes for resuming
                        nodeCount = calculateTaskCount(nodeSelector.selectableNodeCount());
                    }
                    stageNodeList = new ArrayList<>(nodeSelector.selectRandomNodes(nodeCount));
                    checkCondition(stageNodeList.size() == nodeCount, NO_NODES_AVAILABLE, "Snapshot: not enough worker nodes to resume expected number of tasks: " + nodeCount);
                    // Make sure bucketNodeMap uses the same node list
                    bucketNodeMap = nodePartitioningManager.getBucketNodeMap(session, partitioningHandle, dynamicLifespanSchedule, stageNodeList);
                } else {
                    bucketNodeMap = nodePartitioningManager.getBucketNodeMap(session, partitioningHandle, dynamicLifespanSchedule);
                    stageNodeList = new ArrayList<>(nodeScheduler.createNodeSelector(catalogName, keepConsumerOnFeederNodes, feederScheduledNodes).allNodes());
                }
                // verify execution is consistent with planner's decision on dynamic lifespan schedule
                verify(bucketNodeMap.isDynamic() == dynamicLifespanSchedule);
                Collections.shuffle(stageNodeList);
                bucketToPartition = Optional.empty();
            } else {
                // cannot use dynamic lifespan schedule
                verify(!plan.getFragment().getStageExecutionDescriptor().isDynamicLifespanSchedule());
                // remote source requires nodePartitionMap
                NodePartitionMap nodePartitionMap = partitioningCache.apply(plan.getFragment().getPartitioning(), stageTaskCounts == null ? null : stageTaskCounts.get(stageId));
                if (groupedExecutionForStage) {
                    checkState(connectorPartitionHandles.size() == nodePartitionMap.getBucketToPartition().length);
                }
                stageNodeList = nodePartitionMap.getPartitionToNode();
                bucketNodeMap = nodePartitionMap.asBucketNodeMap();
                bucketToPartition = Optional.of(nodePartitionMap.getBucketToPartition());
            }
            stageSchedulers.put(stageId, new FixedSourcePartitionedScheduler(stageExecution, splitSources, plan.getFragment().getStageExecutionDescriptor(), schedulingOrder, stageNodeList, bucketNodeMap, splitBatchSize, getConcurrentLifespansPerNode(session), nodeScheduler.createNodeSelector(catalogName, keepConsumerOnFeederNodes, feederScheduledNodes), connectorPartitionHandles, session, heuristicIndexerManager));
        } else {
            // all sources are remote
            NodePartitionMap nodePartitionMap = partitioningCache.apply(plan.getFragment().getPartitioning(), stageTaskCounts == null ? null : stageTaskCounts.get(stageId));
            List<InternalNode> partitionToNode = nodePartitionMap.getPartitionToNode();
            // todo this should asynchronously wait a standard timeout period before failing
            checkCondition(!partitionToNode.isEmpty(), NO_NODES_AVAILABLE, "No worker nodes available");
            stageSchedulers.put(stageId, new FixedCountScheduler(stageExecution, partitionToNode));
            bucketToPartition = Optional.of(nodePartitionMap.getBucketToPartition());
        }
    }
    ImmutableSet.Builder<SqlStageExecution> childStagesBuilder = ImmutableSet.builder();
    for (StageExecutionPlan subStagePlan : plan.getSubStages()) {
        if (visitedPlanFrags.contains(subStagePlan.getFragment().getId())) {
            continue;
        }
        visitedPlanFrags.add(subStagePlan.getFragment().getId());
        List<SqlStageExecution> subTree = createStages(stageExecution::addExchangeLocations, nextStageId, locationFactory, subStagePlan.withBucketToPartition(bucketToPartition), nodeScheduler, remoteTaskFactory, session, splitBatchSize, partitioningCache, nodePartitioningManager, queryExecutor, schedulerExecutor, failureDetector, nodeTaskMap, stageSchedulers, stageLinkages, isSnapshotEnabled, snapshotManager, stageTaskCounts);
        localStages.addAll(subTree);
        SqlStageExecution childStage = subTree.get(0);
        childStagesBuilder.add(childStage);
        Optional<RemoteSourceNode> parentNode = plan.getFragment().getRemoteSourceNodes().stream().filter(x -> x.getSourceFragmentIds().contains(childStage.getFragment().getId())).findAny();
        checkArgument(parentNode.isPresent(), "Couldn't find parent of a CTE node");
        childStage.setParentId(parentNode.get().getId());
    }
    Set<SqlStageExecution> childStages = childStagesBuilder.build();
    stageExecution.addStateChangeListener(newState -> {
        if (newState.isDone() && newState != StageState.RESCHEDULING) {
            // Snapshot: For "rescheduling", tasks are already cancelled (for resume)
            childStages.forEach(SqlStageExecution::cancel);
        }
    });
    stageLinkages.put(stageId, new StageLinkage(plan.getFragment().getId(), parent, childStages));
    if (partitioningHandle.equals(SCALED_WRITER_DISTRIBUTION)) {
        Supplier<Collection<TaskStatus>> sourceTasksProvider = () -> childStages.stream().map(SqlStageExecution::getAllTasks).flatMap(Collection::stream).map(RemoteTask::getTaskStatus).collect(toList());
        Supplier<Collection<TaskStatus>> writerTasksProvider = () -> stageExecution.getAllTasks().stream().map(RemoteTask::getTaskStatus).collect(toList());
        ScaledWriterScheduler scheduler = new ScaledWriterScheduler(stageExecution, sourceTasksProvider, writerTasksProvider, nodeScheduler.createNodeSelector(null, keepConsumerOnFeederNodes, feederScheduledNodes), schedulerExecutor, getWriterMinSize(session), isSnapshotEnabled, stageTaskCounts != null ? stageTaskCounts.get(stageId) : null);
        whenAllStages(childStages, StageState::isDone).addListener(scheduler::finish, directExecutor());
        stageSchedulers.put(stageId, scheduler);
    }
    return localStages.build();
}

Example 15 with RemoteTask

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

the class SourcePartitionedScheduler method assignSplits.

private Set<RemoteTask> assignSplits(Multimap<InternalNode, Split> splitAssignment, Multimap<InternalNode, Lifespan> noMoreSplitsNotification) {
    ImmutableSet.Builder<RemoteTask> newTasks = ImmutableSet.builder();
    ImmutableSet<InternalNode> nodes = ImmutableSet.<InternalNode>builder().addAll(splitAssignment.keySet()).addAll(noMoreSplitsNotification.keySet()).build();
    for (InternalNode node : nodes) {
        // source partitioned tasks can only receive broadcast data; otherwise it would have a different distribution
        ImmutableMultimap<PlanNodeId, Split> splits = ImmutableMultimap.<PlanNodeId, Split>builder().putAll(partitionedNode, splitAssignment.get(node)).build();
        ImmutableMultimap.Builder<PlanNodeId, Lifespan> noMoreSplits = ImmutableMultimap.builder();
        if (noMoreSplitsNotification.containsKey(node)) {
            noMoreSplits.putAll(partitionedNode, noMoreSplitsNotification.get(node));
        }
        newTasks.addAll(stage.scheduleSplits(node, splits, noMoreSplits.build()));
    }
    return newTasks.build();
}