Examples with ShuffleEnvironment org.apache.flink.runtime.shuffle.ShuffleEnvironment used on opensource projects

Example 1 with ShuffleEnvironment

use of org.apache.flink.runtime.shuffle.ShuffleEnvironment in project flink by apache.

the class TaskExecutorPartitionLifecycleTest method internalTestPartitionRelease.

private void internalTestPartitionRelease(TaskExecutorPartitionTracker partitionTracker, ShuffleEnvironment<?, ?> shuffleEnvironment, CompletableFuture<ResultPartitionID> startTrackingFuture, TestAction testAction) throws Exception {
    final ResultPartitionDeploymentDescriptor taskResultPartitionDescriptor = PartitionTestUtils.createPartitionDeploymentDescriptor(ResultPartitionType.BLOCKING);
    final ExecutionAttemptID eid1 = taskResultPartitionDescriptor.getShuffleDescriptor().getResultPartitionID().getProducerId();
    final TaskDeploymentDescriptor taskDeploymentDescriptor = TaskExecutorSubmissionTest.createTaskDeploymentDescriptor(jobId, "job", eid1, new SerializedValue<>(new ExecutionConfig()), "Sender", 1, 0, 1, 0, new Configuration(), new Configuration(), TestingInvokable.class.getName(), Collections.singletonList(taskResultPartitionDescriptor), Collections.emptyList(), Collections.emptyList(), Collections.emptyList());
    final TaskSlotTable<Task> taskSlotTable = createTaskSlotTable();
    final TaskExecutorLocalStateStoresManager localStateStoresManager = new TaskExecutorLocalStateStoresManager(false, Reference.owned(new File[] { tmp.newFolder() }), Executors.directExecutor());
    final TaskManagerServices taskManagerServices = new TaskManagerServicesBuilder().setTaskSlotTable(taskSlotTable).setTaskStateManager(localStateStoresManager).setShuffleEnvironment(shuffleEnvironment).build();
    final CompletableFuture<Void> taskFinishedFuture = new CompletableFuture<>();
    final OneShotLatch slotOfferedLatch = new OneShotLatch();
    final TestingJobMasterGateway jobMasterGateway = new TestingJobMasterGatewayBuilder().setRegisterTaskManagerFunction((ignoredJobId, ignoredTaskManagerRegistrationInformation) -> CompletableFuture.completedFuture(new JMTMRegistrationSuccess(ResourceID.generate()))).setOfferSlotsFunction((resourceID, slotOffers) -> {
        slotOfferedLatch.trigger();
        return CompletableFuture.completedFuture(slotOffers);
    }).setUpdateTaskExecutionStateFunction(taskExecutionState -> {
        if (taskExecutionState.getExecutionState() == ExecutionState.FINISHED) {
            taskFinishedFuture.complete(null);
        }
        return CompletableFuture.completedFuture(Acknowledge.get());
    }).build();
    final TestingTaskExecutor taskExecutor = createTestingTaskExecutor(taskManagerServices, partitionTracker);
    final CompletableFuture<SlotReport> initialSlotReportFuture = new CompletableFuture<>();
    final TestingResourceManagerGateway testingResourceManagerGateway = new TestingResourceManagerGateway();
    testingResourceManagerGateway.setSendSlotReportFunction(resourceIDInstanceIDSlotReportTuple3 -> {
        initialSlotReportFuture.complete(resourceIDInstanceIDSlotReportTuple3.f2);
        return CompletableFuture.completedFuture(Acknowledge.get());
    });
    testingResourceManagerGateway.setRegisterTaskExecutorFunction(input -> CompletableFuture.completedFuture(new TaskExecutorRegistrationSuccess(new InstanceID(), testingResourceManagerGateway.getOwnResourceId(), new ClusterInformation("blobServerHost", 55555))));
    try {
        taskExecutor.start();
        taskExecutor.waitUntilStarted();
        final TaskExecutorGateway taskExecutorGateway = taskExecutor.getSelfGateway(TaskExecutorGateway.class);
        final String jobMasterAddress = "jm";
        rpc.registerGateway(jobMasterAddress, jobMasterGateway);
        rpc.registerGateway(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway);
        // inform the task manager about the job leader
        taskManagerServices.getJobLeaderService().addJob(jobId, jobMasterAddress);
        jobManagerLeaderRetriever.notifyListener(jobMasterAddress, UUID.randomUUID());
        resourceManagerLeaderRetriever.notifyListener(testingResourceManagerGateway.getAddress(), testingResourceManagerGateway.getFencingToken().toUUID());
        final Optional<SlotStatus> slotStatusOptional = StreamSupport.stream(initialSlotReportFuture.get().spliterator(), false).findAny();
        assertTrue(slotStatusOptional.isPresent());
        final SlotStatus slotStatus = slotStatusOptional.get();
        while (true) {
            try {
                taskExecutorGateway.requestSlot(slotStatus.getSlotID(), jobId, taskDeploymentDescriptor.getAllocationId(), ResourceProfile.ZERO, jobMasterAddress, testingResourceManagerGateway.getFencingToken(), timeout).get();
                break;
            } catch (Exception e) {
                // the proper establishment of the RM connection is tracked
                // asynchronously, so we have to poll here until it went through
                // until then, slot requests will fail with an exception
                Thread.sleep(50);
            }
        }
        TestingInvokable.sync = new BlockerSync();
        // Wait till the slot has been successfully offered before submitting the task.
        // This ensures TM has been successfully registered to JM.
        slotOfferedLatch.await();
        taskExecutorGateway.submitTask(taskDeploymentDescriptor, jobMasterGateway.getFencingToken(), timeout).get();
        TestingInvokable.sync.awaitBlocker();
        // the task is still running => the partition is in in-progress and should be tracked
        assertThat(startTrackingFuture.get(), equalTo(taskResultPartitionDescriptor.getShuffleDescriptor().getResultPartitionID()));
        TestingInvokable.sync.releaseBlocker();
        taskFinishedFuture.get(timeout.getSize(), timeout.getUnit());
        testAction.accept(jobId, taskResultPartitionDescriptor, taskExecutor, taskExecutorGateway);
    } finally {
        RpcUtils.terminateRpcEndpoint(taskExecutor, timeout);
    }
    // the shutdown of the backing shuffle environment releases all partitions
    // the book-keeping is not aware of this
    assertTrue(shuffleEnvironment.getPartitionsOccupyingLocalResources().isEmpty());
}

Example 2 with ShuffleEnvironment

use of org.apache.flink.runtime.shuffle.ShuffleEnvironment in project flink by apache.

the class TaskAsyncCallTest method createTask.

private Task createTask(Class<? extends AbstractInvokable> invokableClass) throws Exception {
    final TestingClassLoaderLease classLoaderHandle = TestingClassLoaderLease.newBuilder().setGetOrResolveClassLoaderFunction((permanentBlobKeys, urls) -> TestingUserCodeClassLoader.newBuilder().setClassLoader(new TestUserCodeClassLoader()).build()).build();
    ResultPartitionConsumableNotifier consumableNotifier = new NoOpResultPartitionConsumableNotifier();
    PartitionProducerStateChecker partitionProducerStateChecker = mock(PartitionProducerStateChecker.class);
    Executor executor = mock(Executor.class);
    TaskMetricGroup taskMetricGroup = UnregisteredMetricGroups.createUnregisteredTaskMetricGroup();
    JobInformation jobInformation = new JobInformation(new JobID(), "Job Name", new SerializedValue<>(new ExecutionConfig()), new Configuration(), Collections.emptyList(), Collections.emptyList());
    TaskInformation taskInformation = new TaskInformation(new JobVertexID(), "Test Task", 1, 1, invokableClass.getName(), new Configuration());
    return new Task(jobInformation, taskInformation, new ExecutionAttemptID(), new AllocationID(), 0, 0, Collections.<ResultPartitionDeploymentDescriptor>emptyList(), Collections.<InputGateDeploymentDescriptor>emptyList(), mock(MemoryManager.class), mock(IOManager.class), shuffleEnvironment, new KvStateService(new KvStateRegistry(), null, null), mock(BroadcastVariableManager.class), new TaskEventDispatcher(), ExternalResourceInfoProvider.NO_EXTERNAL_RESOURCES, new TestTaskStateManager(), mock(TaskManagerActions.class), mock(InputSplitProvider.class), mock(CheckpointResponder.class), new NoOpTaskOperatorEventGateway(), new TestGlobalAggregateManager(), classLoaderHandle, mock(FileCache.class), new TestingTaskManagerRuntimeInfo(), taskMetricGroup, consumableNotifier, partitionProducerStateChecker, executor);
}

Example 3 with ShuffleEnvironment

use of org.apache.flink.runtime.shuffle.ShuffleEnvironment in project flink by apache.

the class TaskSubmissionTestEnvironment method createShuffleEnvironment.

private static ShuffleEnvironment<?, ?> createShuffleEnvironment(ResourceID taskManagerLocation, boolean localCommunication, Configuration configuration, RpcService testingRpcService, boolean mockShuffleEnvironment) throws Exception {
    final ShuffleEnvironment<?, ?> shuffleEnvironment;
    if (mockShuffleEnvironment) {
        shuffleEnvironment = mock(ShuffleEnvironment.class, Mockito.RETURNS_MOCKS);
    } else {
        final InetSocketAddress socketAddress = new InetSocketAddress(InetAddress.getByName(testingRpcService.getAddress()), configuration.getInteger(NettyShuffleEnvironmentOptions.DATA_PORT));
        final NettyConfig nettyConfig = new NettyConfig(socketAddress.getAddress(), socketAddress.getPort(), ConfigurationParserUtils.getPageSize(configuration), ConfigurationParserUtils.getSlot(configuration), configuration);
        shuffleEnvironment = new NettyShuffleEnvironmentBuilder().setTaskManagerLocation(taskManagerLocation).setPartitionRequestInitialBackoff(configuration.getInteger(NettyShuffleEnvironmentOptions.NETWORK_REQUEST_BACKOFF_INITIAL)).setPartitionRequestMaxBackoff(configuration.getInteger(NettyShuffleEnvironmentOptions.NETWORK_REQUEST_BACKOFF_MAX)).setNettyConfig(localCommunication ? null : nettyConfig).build();
        shuffleEnvironment.start();
    }
    return shuffleEnvironment;
}

Example 4 with ShuffleEnvironment

use of org.apache.flink.runtime.shuffle.ShuffleEnvironment in project flink by apache.

the class AbstractUdfStreamOperatorLifecycleTest method testLifeCycleCancel.

@Test
public void testLifeCycleCancel() throws Exception {
    ACTUAL_ORDER_TRACKING.clear();
    Configuration taskManagerConfig = new Configuration();
    StreamConfig cfg = new StreamConfig(new Configuration());
    MockSourceFunction srcFun = new MockSourceFunction();
    cfg.setStreamOperator(new LifecycleTrackingStreamSource<>(srcFun, false));
    cfg.setOperatorID(new OperatorID());
    cfg.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);
    try (ShuffleEnvironment shuffleEnvironment = new NettyShuffleEnvironmentBuilder().build()) {
        Task task = StreamTaskTest.createTask(SourceStreamTask.class, shuffleEnvironment, cfg, taskManagerConfig);
        task.startTaskThread();
        LifecycleTrackingStreamSource.runStarted.await();
        // this should cancel the task even though it is blocked on runFinished
        task.cancelExecution();
        // wait for clean termination
        task.getExecutingThread().join();
        assertEquals(ExecutionState.CANCELED, task.getExecutionState());
        assertEquals(EXPECTED_CALL_ORDER_CANCEL_RUNNING, ACTUAL_ORDER_TRACKING);
    }
}

Example 5 with ShuffleEnvironment

use of org.apache.flink.runtime.shuffle.ShuffleEnvironment in project flink by apache.

the class StreamTaskTest method testStateBackendLoadingAndClosing.

@Test
public void testStateBackendLoadingAndClosing() throws Exception {
    Configuration taskManagerConfig = new Configuration();
    taskManagerConfig.setString(STATE_BACKEND, TestMemoryStateBackendFactory.class.getName());
    StreamConfig cfg = new StreamConfig(new Configuration());
    cfg.setStateKeySerializer(mock(TypeSerializer.class));
    cfg.setOperatorID(new OperatorID(4711L, 42L));
    TestStreamSource<Long, MockSourceFunction> streamSource = new TestStreamSource<>(new MockSourceFunction());
    cfg.setStreamOperator(streamSource);
    cfg.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);
    try (ShuffleEnvironment shuffleEnvironment = new NettyShuffleEnvironmentBuilder().build()) {
        Task task = createTask(StateBackendTestSource.class, shuffleEnvironment, cfg, taskManagerConfig);
        StateBackendTestSource.fail = false;
        task.startTaskThread();
        // wait for clean termination
        task.getExecutingThread().join();
        // ensure that the state backends and stream iterables are closed ...
        verify(TestStreamSource.operatorStateBackend).close();
        verify(TestStreamSource.keyedStateBackend).close();
        verify(TestStreamSource.rawOperatorStateInputs).close();
        verify(TestStreamSource.rawKeyedStateInputs).close();
        // ... and disposed
        verify(TestStreamSource.operatorStateBackend).dispose();
        verify(TestStreamSource.keyedStateBackend).dispose();
        assertEquals(ExecutionState.FINISHED, task.getExecutionState());
    }
}