use of com.facebook.presto.execution.SqlStageExecution in project presto by prestodb.
the class SqlQueryScheduler method cancelStage.
public void cancelStage(StageId stageId) {
try (SetThreadName ignored = new SetThreadName("Query-%s", queryStateMachine.getQueryId())) {
SqlStageExecution sqlStageExecution = stages.get(stageId);
SqlStageExecution stage = requireNonNull(sqlStageExecution, () -> format("Stage %s does not exist", stageId));
stage.cancel();
}
}
use of com.facebook.presto.execution.SqlStageExecution in project presto by prestodb.
the class TestSourcePartitionedScheduler method testBlockCausesFullSchedule.
@Test
public void testBlockCausesFullSchedule() {
NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService);
// Schedule 60 splits - filling up all nodes
SubPlan firstPlan = createPlan();
SqlStageExecution firstStage = createSqlStageExecution(firstPlan, nodeTaskMap);
StageScheduler firstScheduler = getSourcePartitionedScheduler(createFixedSplitSource(60, TestingSplit::createRemoteSplit), firstStage, nodeManager, nodeTaskMap, 200);
ScheduleResult scheduleResult = firstScheduler.schedule();
assertEffectivelyFinished(scheduleResult, firstScheduler);
assertTrue(scheduleResult.getBlocked().isDone());
assertEquals(scheduleResult.getNewTasks().size(), 3);
assertEquals(firstStage.getAllTasks().size(), 3);
for (RemoteTask remoteTask : firstStage.getAllTasks()) {
PartitionedSplitsInfo splitsInfo = remoteTask.getPartitionedSplitsInfo();
assertEquals(splitsInfo.getCount(), 20);
}
// Schedule more splits in another query, which will block since all nodes are full
SubPlan secondPlan = createPlan();
SqlStageExecution secondStage = createSqlStageExecution(secondPlan, nodeTaskMap);
StageScheduler secondScheduler = getSourcePartitionedScheduler(createFixedSplitSource(5, TestingSplit::createRemoteSplit), secondStage, nodeManager, nodeTaskMap, 200);
scheduleResult = secondScheduler.schedule();
assertFalse(scheduleResult.isFinished());
assertTrue(scheduleResult.getBlocked().isDone());
assertEquals(scheduleResult.getNewTasks().size(), 3);
assertEquals(secondStage.getAllTasks().size(), 3);
for (RemoteTask remoteTask : secondStage.getAllTasks()) {
PartitionedSplitsInfo splitsInfo = remoteTask.getPartitionedSplitsInfo();
assertEquals(splitsInfo.getCount(), 0);
}
firstStage.abort();
secondStage.abort();
}
use of com.facebook.presto.execution.SqlStageExecution in project presto by prestodb.
the class TestSourcePartitionedScheduler method testBalancedSplitAssignment.
@Test
public void testBalancedSplitAssignment() {
// use private node manager so we can add a node later
InMemoryNodeManager nodeManager = new InMemoryNodeManager();
nodeManager.addNode(CONNECTOR_ID, new InternalNode("other1", URI.create("http://127.0.0.1:11"), NodeVersion.UNKNOWN, false), new InternalNode("other2", URI.create("http://127.0.0.1:12"), NodeVersion.UNKNOWN, false), new InternalNode("other3", URI.create("http://127.0.0.1:13"), NodeVersion.UNKNOWN, false));
NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService);
// Schedule 15 splits - there are 3 nodes, each node should get 5 splits
SubPlan firstPlan = createPlan();
SqlStageExecution firstStage = createSqlStageExecution(firstPlan, nodeTaskMap);
StageScheduler firstScheduler = getSourcePartitionedScheduler(createFixedSplitSource(15, TestingSplit::createRemoteSplit), firstStage, nodeManager, nodeTaskMap, 200);
ScheduleResult scheduleResult = firstScheduler.schedule();
assertEffectivelyFinished(scheduleResult, firstScheduler);
assertTrue(scheduleResult.getBlocked().isDone());
assertEquals(scheduleResult.getNewTasks().size(), 3);
assertEquals(firstStage.getAllTasks().size(), 3);
for (RemoteTask remoteTask : firstStage.getAllTasks()) {
PartitionedSplitsInfo splitsInfo = remoteTask.getPartitionedSplitsInfo();
assertEquals(splitsInfo.getCount(), 5);
}
// Add new node
InternalNode additionalNode = new InternalNode("other4", URI.create("http://127.0.0.1:14"), NodeVersion.UNKNOWN, false);
nodeManager.addNode(CONNECTOR_ID, additionalNode);
// Schedule 5 splits in another query. Since the new node does not have any splits, all 5 splits are assigned to the new node
SubPlan secondPlan = createPlan();
SqlStageExecution secondStage = createSqlStageExecution(secondPlan, nodeTaskMap);
StageScheduler secondScheduler = getSourcePartitionedScheduler(createFixedSplitSource(5, TestingSplit::createRemoteSplit), secondStage, nodeManager, nodeTaskMap, 200);
scheduleResult = secondScheduler.schedule();
assertEffectivelyFinished(scheduleResult, secondScheduler);
assertTrue(scheduleResult.getBlocked().isDone());
assertEquals(scheduleResult.getNewTasks().size(), 1);
assertEquals(secondStage.getAllTasks().size(), 1);
RemoteTask task = secondStage.getAllTasks().get(0);
assertEquals(task.getPartitionedSplitsInfo().getCount(), 5);
firstStage.abort();
secondStage.abort();
}
use of com.facebook.presto.execution.SqlStageExecution in project presto by prestodb.
the class TestSourcePartitionedScheduler method testScheduleNoSplits.
@Test
public void testScheduleNoSplits() {
SubPlan plan = createPlan();
NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService);
SqlStageExecution stage = createSqlStageExecution(plan, nodeTaskMap);
StageScheduler scheduler = getSourcePartitionedScheduler(createFixedSplitSource(0, TestingSplit::createRemoteSplit), stage, nodeManager, nodeTaskMap, 1);
ScheduleResult scheduleResult = scheduler.schedule();
assertEquals(scheduleResult.getNewTasks().size(), 1);
assertEffectivelyFinished(scheduleResult, scheduler);
stage.abort();
}
use of com.facebook.presto.execution.SqlStageExecution in project presto by prestodb.
the class LegacySqlQueryScheduler method schedule.
private void schedule() {
if (!scheduling.compareAndSet(false, true)) {
// still scheduling the previous batch of stages
return;
}
List<StageExecutionAndScheduler> scheduledStageExecutions = new ArrayList<>();
try (SetThreadName ignored = new SetThreadName("Query-%s", queryStateMachine.getQueryId())) {
Set<StageId> completedStages = new HashSet<>();
List<ExecutionSchedule> sectionExecutionSchedules = new LinkedList<>();
while (!Thread.currentThread().isInterrupted()) {
// remove finished section
sectionExecutionSchedules.removeIf(ExecutionSchedule::isFinished);
// try to pull more section that are ready to be run
List<StreamingPlanSection> sectionsReadyForExecution = getSectionsReadyForExecution();
// all finished
if (sectionsReadyForExecution.isEmpty() && sectionExecutionSchedules.isEmpty()) {
break;
}
List<List<StageExecutionAndScheduler>> sectionStageExecutions = getStageExecutions(sectionsReadyForExecution);
sectionStageExecutions.forEach(scheduledStageExecutions::addAll);
sectionStageExecutions.stream().map(executionInfos -> executionInfos.stream().collect(toImmutableList())).map(executionPolicy::createExecutionSchedule).forEach(sectionExecutionSchedules::add);
while (sectionExecutionSchedules.stream().noneMatch(ExecutionSchedule::isFinished)) {
List<ListenableFuture<?>> blockedStages = new ArrayList<>();
List<StageExecutionAndScheduler> executionsToSchedule = sectionExecutionSchedules.stream().flatMap(schedule -> schedule.getStagesToSchedule().stream()).collect(toImmutableList());
for (StageExecutionAndScheduler stageExecutionAndScheduler : executionsToSchedule) {
SqlStageExecution stageExecution = stageExecutionAndScheduler.getStageExecution();
StageId stageId = stageExecution.getStageExecutionId().getStageId();
stageExecution.beginScheduling();
// perform some scheduling work
ScheduleResult result = stageExecutionAndScheduler.getStageScheduler().schedule();
// Track leaf tasks if partial results are enabled
if (isPartialResultsEnabled(session) && stageExecutionAndScheduler.getStageExecution().getFragment().isLeaf()) {
for (RemoteTask task : result.getNewTasks()) {
partialResultQueryTaskTracker.trackTask(task);
task.addFinalTaskInfoListener(partialResultQueryTaskTracker::recordTaskFinish);
}
}
// modify parent and children based on the results of the scheduling
if (result.isFinished()) {
stageExecution.schedulingComplete();
} else if (!result.getBlocked().isDone()) {
blockedStages.add(result.getBlocked());
}
stageExecutionAndScheduler.getStageLinkage().processScheduleResults(stageExecution.getState(), result.getNewTasks());
schedulerStats.getSplitsScheduledPerIteration().add(result.getSplitsScheduled());
if (result.getBlockedReason().isPresent()) {
switch(result.getBlockedReason().get()) {
case WRITER_SCALING:
// no-op
break;
case WAITING_FOR_SOURCE:
schedulerStats.getWaitingForSource().update(1);
break;
case SPLIT_QUEUES_FULL:
schedulerStats.getSplitQueuesFull().update(1);
break;
case MIXED_SPLIT_QUEUES_FULL_AND_WAITING_FOR_SOURCE:
schedulerStats.getMixedSplitQueuesFullAndWaitingForSource().update(1);
break;
case NO_ACTIVE_DRIVER_GROUP:
schedulerStats.getNoActiveDriverGroup().update(1);
break;
default:
throw new UnsupportedOperationException("Unknown blocked reason: " + result.getBlockedReason().get());
}
}
}
// make sure to update stage linkage at least once per loop to catch async state changes (e.g., partial cancel)
boolean stageFinishedExecution = false;
for (StageExecutionAndScheduler stageExecutionInfo : scheduledStageExecutions) {
SqlStageExecution stageExecution = stageExecutionInfo.getStageExecution();
StageId stageId = stageExecution.getStageExecutionId().getStageId();
if (!completedStages.contains(stageId) && stageExecution.getState().isDone()) {
stageExecutionInfo.getStageLinkage().processScheduleResults(stageExecution.getState(), ImmutableSet.of());
completedStages.add(stageId);
stageFinishedExecution = true;
}
}
// if any stage has just finished execution try to pull more sections for scheduling
if (stageFinishedExecution) {
break;
}
// wait for a state change and then schedule again
if (!blockedStages.isEmpty()) {
try (TimeStat.BlockTimer timer = schedulerStats.getSleepTime().time()) {
tryGetFutureValue(whenAnyComplete(blockedStages), 1, SECONDS);
}
for (ListenableFuture<?> blockedStage : blockedStages) {
blockedStage.cancel(true);
}
}
}
}
for (StageExecutionAndScheduler stageExecutionInfo : scheduledStageExecutions) {
StageExecutionState state = stageExecutionInfo.getStageExecution().getState();
if (state != SCHEDULED && state != RUNNING && !state.isDone()) {
throw new PrestoException(GENERIC_INTERNAL_ERROR, format("Scheduling is complete, but stage execution %s is in state %s", stageExecutionInfo.getStageExecution().getStageExecutionId(), state));
}
}
scheduling.set(false);
// Inform the tracker that task scheduling has completed
partialResultQueryTaskTracker.completeTaskScheduling();
if (!getSectionsReadyForExecution().isEmpty()) {
startScheduling();
}
} catch (Throwable t) {
scheduling.set(false);
queryStateMachine.transitionToFailed(t);
throw t;
} finally {
RuntimeException closeError = new RuntimeException();
for (StageExecutionAndScheduler stageExecutionInfo : scheduledStageExecutions) {
try {
stageExecutionInfo.getStageScheduler().close();
} catch (Throwable t) {
queryStateMachine.transitionToFailed(t);
// Self-suppression not permitted
if (closeError != t) {
closeError.addSuppressed(t);
}
}
}
if (closeError.getSuppressed().length > 0) {
throw closeError;
}
}
}
Aggregations