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

Example 21 with RemoteTask

use of io.trino.execution.RemoteTask in project trino by trinodb.

the class TestHttpRemoteTask method runTest.

private void runTest(FailureScenario failureScenario) throws Exception {
    AtomicLong lastActivityNanos = new AtomicLong(System.nanoTime());
    TestingTaskResource testingTaskResource = new TestingTaskResource(lastActivityNanos, failureScenario);
    HttpRemoteTaskFactory httpRemoteTaskFactory = createHttpRemoteTaskFactory(testingTaskResource);
    RemoteTask remoteTask = createRemoteTask(httpRemoteTaskFactory, ImmutableSet.of());
    testingTaskResource.setInitialTaskInfo(remoteTask.getTaskInfo());
    remoteTask.start();
    waitUntilIdle(lastActivityNanos);
    httpRemoteTaskFactory.stop();
    assertTrue(remoteTask.getTaskStatus().getState().isDone(), format("TaskStatus is not in a done state: %s", remoteTask.getTaskStatus()));
    ErrorCode actualErrorCode = getOnlyElement(remoteTask.getTaskStatus().getFailures()).getErrorCode();
    switch(failureScenario) {
        case TASK_MISMATCH:
        case TASK_MISMATCH_WHEN_VERSION_IS_HIGH:
            assertTrue(remoteTask.getTaskInfo().getTaskStatus().getState().isDone(), format("TaskInfo is not in a done state: %s", remoteTask.getTaskInfo()));
            assertEquals(actualErrorCode, REMOTE_TASK_MISMATCH.toErrorCode());
            break;
        case REJECTED_EXECUTION:
            // for a rejection to occur, the http client must be shutdown, which means we will not be able to ge the final task info
            assertEquals(actualErrorCode, REMOTE_TASK_ERROR.toErrorCode());
            break;
        default:
            throw new UnsupportedOperationException();
    }
}

Example 22 with RemoteTask

use of io.trino.execution.RemoteTask in project trino by trinodb.

the class UniformNodeSelector method computeAssignments.

@Override
public SplitPlacementResult computeAssignments(Set<Split> splits, List<RemoteTask> existingTasks) {
    Multimap<InternalNode, Split> assignment = HashMultimap.create();
    NodeMap nodeMap = this.nodeMap.get().get();
    NodeAssignmentStats assignmentStats = new NodeAssignmentStats(nodeTaskMap, nodeMap, existingTasks);
    ResettableRandomizedIterator<InternalNode> randomCandidates = randomizedNodes(nodeMap, includeCoordinator, ImmutableSet.of());
    Set<InternalNode> blockedExactNodes = new HashSet<>();
    boolean splitWaitingForAnyNode = false;
    // splitsToBeRedistributed becomes true only when splits go through locality-based assignment
    boolean splitsToBeRedistributed = false;
    Set<Split> remainingSplits = new HashSet<>();
    // optimizedLocalScheduling enables prioritized assignment of splits to local nodes when splits contain locality information
    if (optimizedLocalScheduling) {
        for (Split split : splits) {
            if (split.isRemotelyAccessible() && !split.getAddresses().isEmpty()) {
                List<InternalNode> candidateNodes = selectExactNodes(nodeMap, split.getAddresses(), includeCoordinator);
                Optional<InternalNode> chosenNode = candidateNodes.stream().filter(ownerNode -> assignmentStats.getTotalSplitsWeight(ownerNode) < maxSplitsWeightPerNode && assignmentStats.getUnacknowledgedSplitCountForStage(ownerNode) < maxUnacknowledgedSplitsPerTask).min(comparingLong(assignmentStats::getTotalSplitsWeight));
                if (chosenNode.isPresent()) {
                    assignment.put(chosenNode.get(), split);
                    assignmentStats.addAssignedSplit(chosenNode.get(), split.getSplitWeight());
                    splitsToBeRedistributed = true;
                    continue;
                }
            }
            remainingSplits.add(split);
        }
    } else {
        remainingSplits = splits;
    }
    for (Split split : remainingSplits) {
        randomCandidates.reset();
        List<InternalNode> candidateNodes;
        if (!split.isRemotelyAccessible()) {
            candidateNodes = selectExactNodes(nodeMap, split.getAddresses(), includeCoordinator);
        } else {
            candidateNodes = selectNodes(minCandidates, randomCandidates);
        }
        if (candidateNodes.isEmpty()) {
            log.debug("No nodes available to schedule %s. Available nodes %s", split, nodeMap.getNodesByHost().keys());
            throw new TrinoException(NO_NODES_AVAILABLE, "No nodes available to run query");
        }
        InternalNode chosenNode = chooseNodeForSplit(assignmentStats, candidateNodes);
        if (chosenNode == null) {
            long minWeight = Long.MAX_VALUE;
            for (InternalNode node : candidateNodes) {
                long queuedWeight = assignmentStats.getQueuedSplitsWeightForStage(node);
                if (queuedWeight <= minWeight && queuedWeight < maxPendingSplitsWeightPerTask && assignmentStats.getUnacknowledgedSplitCountForStage(node) < maxUnacknowledgedSplitsPerTask) {
                    chosenNode = node;
                    minWeight = queuedWeight;
                }
            }
        }
        if (chosenNode != null) {
            assignment.put(chosenNode, split);
            assignmentStats.addAssignedSplit(chosenNode, split.getSplitWeight());
        } else {
            if (split.isRemotelyAccessible()) {
                splitWaitingForAnyNode = true;
            } else // Exact node set won't matter, if a split is waiting for any node
            if (!splitWaitingForAnyNode) {
                blockedExactNodes.addAll(candidateNodes);
            }
        }
    }
    ListenableFuture<Void> blocked;
    if (splitWaitingForAnyNode) {
        blocked = toWhenHasSplitQueueSpaceFuture(existingTasks, calculateLowWatermark(maxPendingSplitsWeightPerTask));
    } else {
        blocked = toWhenHasSplitQueueSpaceFuture(blockedExactNodes, existingTasks, calculateLowWatermark(maxPendingSplitsWeightPerTask));
    }
    if (splitsToBeRedistributed) {
        equateDistribution(assignment, assignmentStats, nodeMap, includeCoordinator);
    }
    return new SplitPlacementResult(blocked, assignment);
}

Example 23 with RemoteTask

use of io.trino.execution.RemoteTask in project trino by trinodb.

the class NodeScheduler method toWhenHasSplitQueueSpaceFuture.

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

Example 24 with RemoteTask

use of io.trino.execution.RemoteTask in project trino by trinodb.

the class PipelinedStageExecution method failTask.

@Override
public synchronized void failTask(TaskId taskId, Throwable failureCause) {
    RemoteTask task = requireNonNull(tasks.get(taskId.getPartitionId()), () -> "task not found: " + taskId);
    task.fail(failureCause);
    fail(failureCause);
}

Example 25 with RemoteTask

use of io.trino.execution.RemoteTask in project trino by trinodb.

the class PipelinedStageExecution method scheduleTask.

@Override
public synchronized Optional<RemoteTask> scheduleTask(InternalNode node, int partition, Multimap<PlanNodeId, Split> initialSplits, Multimap<PlanNodeId, Lifespan> noMoreSplitsForLifespan) {
    if (stateMachine.getState().isDone()) {
        return Optional.empty();
    }
    checkArgument(!tasks.containsKey(partition), "A task for partition %s already exists", partition);
    OutputBuffers outputBuffers = outputBufferManagers.get(stage.getFragment().getId()).getOutputBuffers();
    Optional<RemoteTask> optionalTask = stage.createTask(node, partition, attempt, bucketToPartition, outputBuffers, initialSplits, ImmutableMultimap.of(), ImmutableSet.of());
    if (optionalTask.isEmpty()) {
        return Optional.empty();
    }
    RemoteTask task = optionalTask.get();
    tasks.put(partition, task);
    ImmutableMultimap.Builder<PlanNodeId, Split> exchangeSplits = ImmutableMultimap.builder();
    sourceTasks.forEach((fragmentId, sourceTask) -> {
        TaskStatus status = sourceTask.getTaskStatus();
        if (status.getState() != TaskState.FINISHED) {
            PlanNodeId planNodeId = exchangeSources.get(fragmentId).getId();
            exchangeSplits.put(planNodeId, createExchangeSplit(sourceTask, task));
        }
    });
    allTasks.add(task.getTaskId());
    task.addSplits(exchangeSplits.build());
    noMoreSplitsForLifespan.forEach(task::noMoreSplits);
    completeSources.forEach(task::noMoreSplits);
    task.addStateChangeListener(this::updateTaskStatus);
    task.addStateChangeListener(this::updateCompletedDriverGroups);
    task.start();
    taskLifecycleListener.taskCreated(stage.getFragment().getId(), task);
    // update output buffers
    OutputBufferId outputBufferId = new OutputBufferId(task.getTaskId().getPartitionId());
    updateSourceTasksOutputBuffers(outputBufferManager -> outputBufferManager.addOutputBuffer(outputBufferId));
    return Optional.of(task);
}