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

Example 21 with RemoteTask

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

the class TestNodeScheduler method testMaxSplitsPerNode.

@Test
public void testMaxSplitsPerNode() {
    setUpNodes();
    InternalNode newNode = new InternalNode("other4", URI.create("http://10.0.0.1:14"), NodeVersion.UNKNOWN, false);
    nodeManager.addNode(CONNECTOR_ID, newNode);
    ImmutableList.Builder<Split> initialSplits = ImmutableList.builder();
    for (int i = 0; i < 10; i++) {
        initialSplits.add(new Split(CONNECTOR_ID, new TestSplitRemote(), Lifespan.taskWide()));
    }
    MockRemoteTaskFactory remoteTaskFactory = new MockRemoteTaskFactory(remoteTaskExecutor, remoteTaskScheduledExecutor);
    // Max out number of splits on node
    TaskId taskId1 = new TaskId("test", 1, 1);
    RemoteTask remoteTask1 = remoteTaskFactory.createTableScanTask(taskId1, newNode, initialSplits.build(), nodeTaskMap.createPartitionedSplitCountTracker(newNode, taskId1));
    nodeTaskMap.addTask(newNode, remoteTask1);
    TaskId taskId2 = new TaskId("test", 1, 2);
    RemoteTask remoteTask2 = remoteTaskFactory.createTableScanTask(taskId2, newNode, initialSplits.build(), nodeTaskMap.createPartitionedSplitCountTracker(newNode, taskId2));
    nodeTaskMap.addTask(newNode, remoteTask2);
    Set<Split> splits = new HashSet<>();
    for (int i = 0; i < 5; i++) {
        splits.add(new Split(CONNECTOR_ID, new TestSplitRemote(), Lifespan.taskWide()));
    }
    Multimap<InternalNode, Split> assignments = nodeSelector.computeAssignments(splits, ImmutableList.copyOf(taskMap.values()), Optional.empty()).getAssignments();
    // no split should be assigned to the newNode, as it already has maxNodeSplits assigned to it
    assertFalse(assignments.keySet().contains(newNode));
    remoteTask1.abort();
    remoteTask2.abort();
    assertEquals(nodeTaskMap.getPartitionedSplitsOnNode(newNode), 0);
}

Example 22 with RemoteTask

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

the class TestNodeScheduler method testEmptyAssignmentWithFullNodes.

@Test
public void testEmptyAssignmentWithFullNodes() {
    InternalNode node1 = new InternalNode("node1", URI.create("http://10.0.0.1:11"), NodeVersion.UNKNOWN, false);
    nodeManager.addNode(CONNECTOR_ID, node1);
    InternalNode node2 = new InternalNode("node2", URI.create("http://10.0.0.1:12"), NodeVersion.UNKNOWN, false);
    nodeManager.addNode(CONNECTOR_ID, node2);
    Set<Split> splits = new LinkedHashSet<>();
    // 20 splits with node1 as local node to be assigned in the first iteration of computeAssignments
    for (int i = 0; i < (20 + 10 + 5) * 2; i++) {
        splits.add(new Split(CONNECTOR_ID, new TestSplitLocal(), Lifespan.taskWide()));
    }
    // computeAssignments just returns a mapping of nodes with splits to be assigned, it does not assign splits
    Multimap<InternalNode, Split> assignments1 = nodeSelector.computeAssignments(splits, ImmutableList.copyOf(taskMap.values()), Optional.empty()).getAssignments();
    assertEquals(assignments1.size(), 40);
    assertEquals(assignments1.keySet().size(), 2);
    assertEquals(assignments1.get(node1).size(), 20);
    assertEquals(assignments1.get(node2).size(), 20);
    MockRemoteTaskFactory remoteTaskFactory = new MockRemoteTaskFactory(remoteTaskExecutor, remoteTaskScheduledExecutor);
    int task = 0;
    for (InternalNode node : assignments1.keySet()) {
        TaskId taskId = new TaskId("test", 1, task);
        task++;
        MockRemoteTaskFactory.MockRemoteTask remoteTask = remoteTaskFactory.createTableScanTask(taskId, node, ImmutableList.copyOf(assignments1.get(node)), nodeTaskMap.createPartitionedSplitCountTracker(node, taskId));
        remoteTask.startSplits(20);
        nodeTaskMap.addTask(node, remoteTask);
        taskMap.put(node, remoteTask);
    }
    Set<Split> unassignedSplits = Sets.difference(splits, new HashSet<>(assignments1.values()));
    assertEquals(unassignedSplits.size(), 30);
    Multimap<InternalNode, Split> assignments2 = nodeSelector.computeAssignments(unassignedSplits, ImmutableList.copyOf(taskMap.values()), Optional.empty()).getAssignments();
    for (InternalNode node : assignments2.keySet()) {
        RemoteTask remoteTask = taskMap.get(node);
        remoteTask.addSplits(ImmutableMultimap.<PlanNodeId, Split>builder().putAll(new PlanNodeId("sourceId"), assignments2.get(node)).build());
    }
    unassignedSplits = Sets.difference(unassignedSplits, new HashSet<>(assignments2.values()));
    assertEquals(unassignedSplits.size(), 10);
    Multimap<InternalNode, Split> assignments3 = nodeSelector.computeAssignments(unassignedSplits, ImmutableList.copyOf(taskMap.values()), Optional.empty()).getAssignments();
    assertTrue(assignments3.isEmpty());
}

Example 23 with RemoteTask

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

the class TestSourcePartitionedScheduler method testBalancedSplitAssignment.

@Test
public void testBalancedSplitAssignment() {
    // use private node manager so we can add a node later
    InMemoryNodeManager memoryNodeManager = new InMemoryNodeManager();
    memoryNodeManager.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
    StageExecutionPlan firstPlan = createPlan(createFixedSplitSource(15, TestingSplit::createRemoteSplit));
    SqlStageExecution firstStage = createSqlStageExecution(firstPlan, nodeTaskMap);
    StageScheduler firstScheduler = getSourcePartitionedScheduler(firstPlan, firstStage, memoryNodeManager, 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()) {
        assertEquals(remoteTask.getPartitionedSplitCount(), 5);
    }
    // Add new node
    InternalNode additionalNode = new InternalNode("other4", URI.create("http://127.0.0.1:14"), NodeVersion.UNKNOWN, false);
    memoryNodeManager.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
    StageExecutionPlan secondPlan = createPlan(createFixedSplitSource(5, TestingSplit::createRemoteSplit));
    SqlStageExecution secondStage = createSqlStageExecution(secondPlan, nodeTaskMap);
    StageScheduler secondScheduler = getSourcePartitionedScheduler(secondPlan, secondStage, memoryNodeManager, 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.getPartitionedSplitCount(), 5);
    firstStage.abort();
    secondStage.abort();
}

Example 24 with RemoteTask

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

the class TestSourcePartitionedScheduler method testScheduleSplitsBatched.

@Test
public void testScheduleSplitsBatched() {
    StageExecutionPlan plan = createPlan(createFixedSplitSource(60, TestingSplit::createRemoteSplit));
    NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService);
    SqlStageExecution stage = createSqlStageExecution(plan, nodeTaskMap);
    StageScheduler scheduler = getSourcePartitionedScheduler(plan, stage, nodeManager, nodeTaskMap, 7);
    for (int i = 0; i <= (60 / 7); i++) {
        ScheduleResult scheduleResult = scheduler.schedule();
        // finishes when last split is fetched
        if (i == (60 / 7)) {
            assertEffectivelyFinished(scheduleResult, scheduler);
        } else {
            assertFalse(scheduleResult.isFinished());
        }
        // never blocks
        assertTrue(scheduleResult.getBlocked().isDone());
        // first three splits create new tasks
        assertEquals(scheduleResult.getNewTasks().size(), i == 0 ? 3 : 0);
        assertEquals(stage.getAllTasks().size(), 3);
        assertPartitionedSplitCount(stage, min((i + 1) * 7, 60));
    }
    for (RemoteTask remoteTask : stage.getAllTasks()) {
        assertEquals(remoteTask.getPartitionedSplitCount(), 20);
    }
    stage.abort();
}

Example 25 with RemoteTask

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

the class TestSourcePartitionedScheduler method testBlockCausesFullSchedule.

@Test
public void testBlockCausesFullSchedule() {
    NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService);
    // Schedule 60 splits - filling up all nodes
    StageExecutionPlan firstPlan = createPlan(createFixedSplitSource(60, TestingSplit::createRemoteSplit));
    SqlStageExecution firstStage = createSqlStageExecution(firstPlan, nodeTaskMap);
    StageScheduler firstScheduler = getSourcePartitionedScheduler(firstPlan, 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()) {
        assertEquals(remoteTask.getPartitionedSplitCount(), 20);
    }
    // Schedule more splits in another query, which will block since all nodes are full
    StageExecutionPlan secondPlan = createPlan(createFixedSplitSource(5, TestingSplit::createRemoteSplit));
    SqlStageExecution secondStage = createSqlStageExecution(secondPlan, nodeTaskMap);
    StageScheduler secondScheduler = getSourcePartitionedScheduler(secondPlan, 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()) {
        assertEquals(remoteTask.getPartitionedSplitCount(), 0);
    }
    firstStage.abort();
    secondStage.abort();
}