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

Example 1 with StageId

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

the class TestTaskSnapshotManager method TestLoadBacktrack.

@Test
public void TestLoadBacktrack() throws Exception {
    queryId = new QueryId("loadbacktrack");
    StageId stageId = new StageId(queryId, 0);
    TaskId taskId = new TaskId(stageId, 0);
    TaskSnapshotManager snapshotManager = new TaskSnapshotManager(taskId, 0, snapshotUtils);
    snapshotUtils.getOrCreateQuerySnapshotManager(queryId, TEST_SNAPSHOT_SESSION);
    // Save operator state
    MockState state = new MockState("state");
    SnapshotStateId stateId = SnapshotStateId.forOperator(1L, taskId, 3, 4, 5);
    snapshotManager.storeState(stateId, state, 0);
    snapshotManager.setTotalComponents(1);
    snapshotManager.succeededToCapture(stateId);
    // Make snapshot manager think that snapshot #1 is complete
    QuerySnapshotManager querySnapshotManager = new QuerySnapshotManager(queryId, snapshotUtils, TEST_SNAPSHOT_SESSION);
    querySnapshotManager.addNewTask(taskId);
    querySnapshotManager.snapshotInitiated(1L);
    querySnapshotManager.updateQueryCapture(taskId, Collections.singletonMap(1L, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL)));
    querySnapshotManager.getResumeSnapshotId();
    SnapshotStateId newStateId = stateId.withSnapshotId(2);
    Optional<Object> loadedState = snapshotManager.loadState(newStateId);
    // Should return state saved for snapshot #1
    assertTrue(loadedState.isPresent());
    Assert.assertEquals(((MockState) loadedState.get()).getState(), state.getState());
}

Example 2 with StageId

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

the class TestSourcePartitionedScheduler method createSqlStageExecution.

private SqlStageExecution createSqlStageExecution(StageExecutionPlan tableScanPlan, NodeTaskMap nodeTaskMap) {
    StageId stageId = new StageId(new QueryId("query"), 0);
    SqlStageExecution stage = SqlStageExecution.createSqlStageExecution(stageId, locationFactory.createStageLocation(stageId), tableScanPlan.getFragment(), tableScanPlan.getTables(), new MockRemoteTaskFactory(queryExecutor, scheduledExecutor), TEST_SESSION, true, nodeTaskMap, queryExecutor, new NoOpFailureDetector(), new SplitSchedulerStats(), new DynamicFilterService(new LocalStateStoreProvider(seedStoreManager)), new QuerySnapshotManager(stageId.getQueryId(), NOOP_SNAPSHOT_UTILS, TEST_SESSION));
    stage.setOutputBuffers(createInitialEmptyOutputBuffers(PARTITIONED).withBuffer(OUT, 0).withNoMoreBufferIds());
    return stage;
}

Example 3 with StageId

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

the class TestHiveIntegrationSmokeTest method testRuseExchangeSplitsGroupNotMatchingBetweenProducerConsumer.

@Test
public void testRuseExchangeSplitsGroupNotMatchingBetweenProducerConsumer() {
    setUpNodes();
    NodeTaskMap nodeTasks = new NodeTaskMap(new FinalizerService());
    StageId stageId = new StageId(new QueryId("query"), 0);
    UUID uuid = UUID.randomUUID();
    PlanFragment testFragmentProducer = createTableScanPlanFragment("build", ReuseExchangeOperator.STRATEGY.REUSE_STRATEGY_PRODUCER, uuid, 1);
    PlanNodeId tableScanNodeId = new PlanNodeId("plan_id");
    StageExecutionPlan producerStageExecutionPlan = new StageExecutionPlan(testFragmentProducer, ImmutableMap.of(tableScanNodeId, new ConnectorAwareSplitSource(CONNECTOR_ID, createFixedSplitSource(0, TestingSplit::createRemoteSplit))), ImmutableList.of(), ImmutableMap.of(tableScanNodeId, new TableInfo(new QualifiedObjectName("test", TEST_SCHEMA, "test"), TupleDomain.all())));
    SqlStageExecution producerStage = createSqlStageExecution(stageId, new TestSqlTaskManager.MockLocationFactory().createStageLocation(stageId), producerStageExecutionPlan.getFragment(), producerStageExecutionPlan.getTables(), new MockRemoteTaskFactory(remoteTaskExecutor, remoteTaskScheduledExecutor), TEST_SESSION_REUSE, true, nodeTasks, remoteTaskExecutor, new NoOpFailureDetector(), new SplitSchedulerStats(), new DynamicFilterService(new LocalStateStoreProvider(new SeedStoreManager(new FileSystemClientManager()))), new QuerySnapshotManager(stageId.getQueryId(), NOOP_SNAPSHOT_UTILS, TEST_SESSION));
    Set<Split> producerSplits = createAndGetSplits(10);
    Multimap<InternalNode, Split> producerAssignment = nodeSelector.computeAssignments(producerSplits, ImmutableList.copyOf(taskMap.values()), Optional.of(producerStage)).getAssignments();
    PlanFragment testFragmentConsumer = createTableScanPlanFragment("build", ReuseExchangeOperator.STRATEGY.REUSE_STRATEGY_CONSUMER, uuid, 1);
    StageExecutionPlan consumerStageExecutionPlan = new StageExecutionPlan(testFragmentConsumer, ImmutableMap.of(tableScanNodeId, new ConnectorAwareSplitSource(CONNECTOR_ID, createFixedSplitSource(0, TestingSplit::createRemoteSplit))), ImmutableList.of(), ImmutableMap.of(tableScanNodeId, new TableInfo(new QualifiedObjectName("test", TEST_SCHEMA, "test"), TupleDomain.all())));
    SqlStageExecution stage = createSqlStageExecution(stageId, new TestSqlTaskManager.MockLocationFactory().createStageLocation(stageId), consumerStageExecutionPlan.getFragment(), consumerStageExecutionPlan.getTables(), new MockRemoteTaskFactory(remoteTaskExecutor, remoteTaskScheduledExecutor), TEST_SESSION_REUSE, true, nodeTasks, remoteTaskExecutor, new NoOpFailureDetector(), new SplitSchedulerStats(), new DynamicFilterService(new LocalStateStoreProvider(new SeedStoreManager(new FileSystemClientManager()))), new QuerySnapshotManager(stageId.getQueryId(), NOOP_SNAPSHOT_UTILS, TEST_SESSION));
    Set<Split> consumerSplits = createAndGetSplits(50);
    try {
        Multimap<InternalNode, Split> consumerAssignment = nodeSelector.computeAssignments(consumerSplits, ImmutableList.copyOf(taskMap.values()), Optional.of(stage)).getAssignments();
    } catch (PrestoException e) {
        assertEquals("Producer & consumer splits are not same", e.getMessage());
        return;
    }
    assertEquals(false, true);
}

Example 4 with StageId

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

the class SqlQueryScheduler method whenAllStages.

private static ListenableFuture<?> whenAllStages(Collection<SqlStageExecution> stages, Predicate<StageState> predicate) {
    checkArgument(!stages.isEmpty(), "stages is empty");
    Set<StageId> stageIds = newConcurrentHashSet(stages.stream().map(SqlStageExecution::getStageId).collect(toSet()));
    SettableFuture<?> future = SettableFuture.create();
    for (SqlStageExecution stage : stages) {
        stage.addStateChangeListener(state -> {
            if (predicate.test(state) && stageIds.remove(stage.getStageId()) && stageIds.isEmpty()) {
                future.set(null);
            }
        });
    }
    return future;
}

Example 5 with StageId

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

the class SqlQueryScheduler method schedule.

private void schedule() {
    try (SetThreadName ignored = new SetThreadName("Query-%s", queryStateMachine.getQueryId())) {
        Set<StageId> completedStages = new HashSet<>();
        ExecutionSchedule executionSchedule = executionPolicy.createExecutionSchedule(stages.values());
        while (!executionSchedule.isFinished()) {
            List<ListenableFuture<?>> blockedStages = new ArrayList<>();
            for (SqlStageExecution stage : executionSchedule.getStagesToSchedule()) {
                if (isReuseTableScanEnabled(session) && !SqlStageExecution.getReuseTableScanMappingIdStatus(stage.getStateMachine())) {
                    continue;
                }
                stage.beginScheduling();
                // and then let it schedule 10% of splits.
                if (queryStateMachine.isThrottlingEnabled() && !canScheduleMoreSplits()) {
                    try {
                        SECONDS.sleep(THROTTLE_SLEEP_TIMER[currentTimerLevel]);
                    } catch (InterruptedException e) {
                        throw new PrestoException(GENERIC_INTERNAL_ERROR, "interrupted while sleeping");
                    }
                    currentTimerLevel = Math.min(currentTimerLevel + 1, THROTTLE_SLEEP_TIMER.length - 1);
                    stage.setThrottledSchedule(true);
                } else {
                    stage.setThrottledSchedule(false);
                    currentTimerLevel = 0;
                }
                // perform some scheduling work
                /* Get groupSize specification from the ResourceGroupManager */
                int maxSplitGroupSize = getOptimalSmallSplitGroupSize();
                ScheduleResult result = stageSchedulers.get(stage.getStageId()).schedule(maxSplitGroupSize);
                // modify parent and children based on the results of the scheduling
                if (result.isFinished()) {
                    stage.schedulingComplete();
                } else if (!result.getBlocked().isDone()) {
                    blockedStages.add(result.getBlocked());
                }
                stageLinkages.get(stage.getStageId()).processScheduleResults(stage.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:
                        case NO_ACTIVE_DRIVER_GROUP:
                            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)
            for (SqlStageExecution stage : stages.values()) {
                if (!completedStages.contains(stage.getStageId()) && stage.getState().isDone()) {
                    stageLinkages.get(stage.getStageId()).processScheduleResults(stage.getState(), ImmutableSet.of());
                    completedStages.add(stage.getStageId());
                }
            }
            // 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 (SqlStageExecution stage : stages.values()) {
            StageState state = stage.getState();
            // Snapshot: if state is resumable_failure, then state of stage and query will change soon again. Don't treat as an error.
            if (state != SCHEDULED && state != RUNNING && !state.isDone() && state != RESUMABLE_FAILURE) {
                throw new PrestoException(GENERIC_INTERNAL_ERROR, format("Scheduling is complete, but stage %s is in state %s", stage.getStageId(), state));
            }
        }
    } catch (Throwable t) {
        queryStateMachine.transitionToFailed(t);
        throw t;
    } finally {
        RuntimeException closeError = new RuntimeException();
        for (StageScheduler scheduler : stageSchedulers.values()) {
            try {
                // Snapshot: when trying to reschedule, then don't close the scheduler (and more importantly, split sources in it)
                QueryState state = queryStateMachine.getQueryState();
                if (state != QueryState.RESCHEDULING && state != QueryState.RESUMING) {
                    scheduler.close();
                }
            } catch (Throwable t) {
                queryStateMachine.transitionToFailed(t);
                // Self-suppression not permitted
                if (closeError != t) {
                    closeError.addSuppressed(t);
                }
            }
        }
        // Snpashot: if resuming, notify the new scheduler so it can start scheduling new stages
        schedulingFuture.set(null);
        if (closeError.getSuppressed().length > 0) {
            throw closeError;
        }
    }
}