Examples with TestingTaskExecutorGatewayBuilder org.apache.flink.runtime.taskexecutor.TestingTaskExecutorGatewayBuilder used on opensource projects

Example 56 with TestingTaskExecutorGatewayBuilder

use of org.apache.flink.runtime.taskexecutor.TestingTaskExecutorGatewayBuilder in project flink by apache.

the class DeclarativeSlotManagerTest method registerTaskExecutorWithTwoSlots.

private void registerTaskExecutorWithTwoSlots(DeclarativeSlotManager slotManager, CompletableFuture<JobID> firstRequestSlotFuture) {
    final TestingTaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().setRequestSlotFunction(slotIDJobIDAllocationIDStringResourceManagerIdTuple6 -> {
        firstRequestSlotFuture.complete(slotIDJobIDAllocationIDStringResourceManagerIdTuple6.f1);
        return CompletableFuture.completedFuture(Acknowledge.get());
    }).createTestingTaskExecutorGateway();
    final TaskExecutorConnection firstTaskExecutorConnection = createTaskExecutorConnection(taskExecutorGateway);
    final SlotReport firstSlotReport = createSlotReport(firstTaskExecutorConnection.getResourceID(), 2);
    slotManager.registerTaskManager(firstTaskExecutorConnection, firstSlotReport, ResourceProfile.ANY, ResourceProfile.ANY);
}

Example 57 with TestingTaskExecutorGatewayBuilder

use of org.apache.flink.runtime.taskexecutor.TestingTaskExecutorGatewayBuilder in project flink by apache.

the class DeclarativeSlotManagerTest method testTaskManagerUnregistration.

/**
 * Tests that un-registration of task managers will free and remove all registered slots.
 */
@Test
public void testTaskManagerUnregistration() throws Exception {
    final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().setRequestSlotFunction(tuple6 -> new CompletableFuture<>()).createTestingTaskExecutorGateway();
    final TaskExecutorConnection taskManagerConnection = createTaskExecutorConnection(taskExecutorGateway);
    final ResourceID resourceId = taskManagerConnection.getResourceID();
    final SlotID slotId1 = new SlotID(resourceId, 0);
    final SlotID slotId2 = new SlotID(resourceId, 1);
    final SlotReport slotReport = new SlotReport(Arrays.asList(createAllocatedSlotStatus(slotId1), createFreeSlotStatus(slotId2)));
    final ResourceRequirements resourceRequirements = createResourceRequirementsForSingleSlot();
    final DefaultSlotTracker slotTracker = new DefaultSlotTracker();
    try (DeclarativeSlotManager slotManager = createDeclarativeSlotManagerBuilder().setSlotTracker(slotTracker).buildAndStartWithDirectExec()) {
        slotManager.registerTaskManager(taskManagerConnection, slotReport, ResourceProfile.ANY, ResourceProfile.ANY);
        assertEquals("The number registered slots does not equal the expected number.", 2, slotManager.getNumberRegisteredSlots());
        slotManager.processResourceRequirements(resourceRequirements);
        slotManager.unregisterTaskManager(taskManagerConnection.getInstanceID(), TEST_EXCEPTION);
        assertEquals(0, slotManager.getNumberRegisteredSlots());
    }
}

Example 58 with TestingTaskExecutorGatewayBuilder

use of org.apache.flink.runtime.taskexecutor.TestingTaskExecutorGatewayBuilder in project flink by apache.

the class FineGrainedSlotManagerTest method testSlotAllocationAccordingToStrategyResult.

// ---------------------------------------------------------------------------------------------
// Handle result from ResourceAllocationStrategy
// ---------------------------------------------------------------------------------------------
@Test
public void testSlotAllocationAccordingToStrategyResult() throws Exception {
    final CompletableFuture<Tuple6<SlotID, JobID, AllocationID, ResourceProfile, String, ResourceManagerId>> requestSlotFuture = new CompletableFuture<>();
    final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().setRequestSlotFunction(tuple6 -> {
        requestSlotFuture.complete(tuple6);
        return CompletableFuture.completedFuture(Acknowledge.get());
    }).createTestingTaskExecutorGateway();
    final TaskExecutorConnection taskManagerConnection = new TaskExecutorConnection(ResourceID.generate(), taskExecutorGateway);
    final JobID jobId = new JobID();
    final SlotReport slotReport = new SlotReport();
    new Context() {

        {
            resourceAllocationStrategyBuilder.setTryFulfillRequirementsFunction(((jobIDCollectionMap, taskManagerResourceInfoProvider) -> ResourceAllocationResult.builder().addAllocationOnRegisteredResource(jobId, taskManagerConnection.getInstanceID(), DEFAULT_SLOT_RESOURCE_PROFILE).build()));
            runTest(() -> {
                runInMainThread(() -> {
                    getSlotManager().registerTaskManager(taskManagerConnection, slotReport, DEFAULT_TOTAL_RESOURCE_PROFILE, DEFAULT_SLOT_RESOURCE_PROFILE);
                    getSlotManager().processResourceRequirements(createResourceRequirements(jobId, 1));
                });
                final Tuple6<SlotID, JobID, AllocationID, ResourceProfile, String, ResourceManagerId> requestSlot = assertFutureCompleteAndReturn(requestSlotFuture);
                assertEquals(jobId, requestSlot.f1);
                assertEquals(DEFAULT_SLOT_RESOURCE_PROFILE, requestSlot.f3);
            });
        }
    };
}

Example 59 with TestingTaskExecutorGatewayBuilder

use of org.apache.flink.runtime.taskexecutor.TestingTaskExecutorGatewayBuilder in project flink by apache.

the class FineGrainedSlotManagerTest method testSlotAllocationForPendingTaskManagerWillBeRespected.

@Test
public void testSlotAllocationForPendingTaskManagerWillBeRespected() throws Exception {
    final JobID jobId = new JobID();
    final CompletableFuture<Void> requestResourceFuture = new CompletableFuture<>();
    final PendingTaskManager pendingTaskManager = new PendingTaskManager(DEFAULT_TOTAL_RESOURCE_PROFILE, DEFAULT_NUM_SLOTS_PER_WORKER);
    final CompletableFuture<Tuple6<SlotID, JobID, AllocationID, ResourceProfile, String, ResourceManagerId>> requestSlotFuture = new CompletableFuture<>();
    final TaskExecutorGateway taskExecutorGateway = new TestingTaskExecutorGatewayBuilder().setRequestSlotFunction(tuple6 -> {
        requestSlotFuture.complete(tuple6);
        return CompletableFuture.completedFuture(Acknowledge.get());
    }).createTestingTaskExecutorGateway();
    final TaskExecutorConnection taskManagerConnection = new TaskExecutorConnection(ResourceID.generate(), taskExecutorGateway);
    new Context() {

        {
            resourceAllocationStrategyBuilder.setTryFulfillRequirementsFunction(((jobIDCollectionMap, taskManagerResourceInfoProvider) -> ResourceAllocationResult.builder().addPendingTaskManagerAllocate(pendingTaskManager).addAllocationOnPendingResource(jobId, pendingTaskManager.getPendingTaskManagerId(), DEFAULT_SLOT_RESOURCE_PROFILE).build()));
            resourceActionsBuilder.setAllocateResourceConsumer(ignored -> requestResourceFuture.complete(null));
            runTest(() -> {
                runInMainThread(() -> getSlotManager().processResourceRequirements(createResourceRequirements(jobId, 1)));
                assertFutureCompleteAndReturn(requestResourceFuture);
                runInMainThread(() -> getSlotManager().registerTaskManager(taskManagerConnection, new SlotReport(), DEFAULT_TOTAL_RESOURCE_PROFILE, DEFAULT_SLOT_RESOURCE_PROFILE));
                final Tuple6<SlotID, JobID, AllocationID, ResourceProfile, String, ResourceManagerId> requestSlot = assertFutureCompleteAndReturn(requestSlotFuture);
                assertEquals(jobId, requestSlot.f1);
                assertEquals(DEFAULT_SLOT_RESOURCE_PROFILE, requestSlot.f3);
            });
        }
    };
}

Example 60 with TestingTaskExecutorGatewayBuilder

use of org.apache.flink.runtime.taskexecutor.TestingTaskExecutorGatewayBuilder in project flink by apache.

the class DefaultSchedulerBatchSchedulingTest method testSchedulingOfJobWithFewerSlotsThanParallelism.

/**
 * Tests that a batch job can be executed with fewer slots than its parallelism. See FLINK-13187
 * for more information.
 */
@Test
public void testSchedulingOfJobWithFewerSlotsThanParallelism() throws Exception {
    final int parallelism = 5;
    final Time batchSlotTimeout = Time.milliseconds(5L);
    final JobGraph jobGraph = createBatchJobGraph(parallelism);
    try (final SlotPool slotPool = createSlotPool(mainThreadExecutor, batchSlotTimeout)) {
        final ArrayBlockingQueue<ExecutionAttemptID> submittedTasksQueue = new ArrayBlockingQueue<>(parallelism);
        TestingTaskExecutorGateway testingTaskExecutorGateway = new TestingTaskExecutorGatewayBuilder().setSubmitTaskConsumer((tdd, ignored) -> {
            submittedTasksQueue.offer(tdd.getExecutionAttemptId());
            return CompletableFuture.completedFuture(Acknowledge.get());
        }).createTestingTaskExecutorGateway();
        final PhysicalSlotProvider slotProvider = new PhysicalSlotProviderImpl(LocationPreferenceSlotSelectionStrategy.createDefault(), slotPool);
        final GloballyTerminalJobStatusListener jobStatusListener = new GloballyTerminalJobStatusListener();
        final SchedulerNG scheduler = createScheduler(jobGraph, mainThreadExecutor, slotProvider, batchSlotTimeout, jobStatusListener);
        CompletableFuture.runAsync(scheduler::startScheduling, mainThreadExecutor).join();
        // register a single slot at the slot pool
        SlotPoolUtils.offerSlots(slotPool, mainThreadExecutor, Collections.singletonList(ResourceProfile.ANY), new RpcTaskManagerGateway(testingTaskExecutorGateway, JobMasterId.generate()));
        // wait until the batch slot timeout has been reached
        Thread.sleep(batchSlotTimeout.toMilliseconds());
        final CompletableFuture<JobStatus> terminationFuture = jobStatusListener.getTerminationFuture();
        for (int i = 0; i < parallelism; i++) {
            final CompletableFuture<ExecutionAttemptID> submittedTaskFuture = CompletableFuture.supplyAsync(CheckedSupplier.unchecked(submittedTasksQueue::take));
            // wait until one of them is completed
            CompletableFuture.anyOf(submittedTaskFuture, terminationFuture).join();
            if (submittedTaskFuture.isDone()) {
                finishExecution(submittedTaskFuture.get(), scheduler, mainThreadExecutor);
            } else {
                fail(String.format("Job reached a globally terminal state %s before all executions were finished.", terminationFuture.get()));
            }
        }
        assertThat(terminationFuture.get(), is(JobStatus.FINISHED));
    }
}