Examples with RemoteEnvironment org.apache.beam.runners.fnexecution.environment.RemoteEnvironment used on opensource projects

Example 6 with RemoteEnvironment

use of org.apache.beam.runners.fnexecution.environment.RemoteEnvironment in project beam by apache.

the class DefaultJobBundleFactoryTest method createsMultipleEnvironmentsWithSdkWorkerParallelism.

@Test
public void createsMultipleEnvironmentsWithSdkWorkerParallelism() throws Exception {
    ServerFactory serverFactory = ServerFactory.createDefault();
    Environment environmentA = Environment.newBuilder().setUrn("env:urn:a").setPayload(ByteString.copyFrom(new byte[1])).build();
    EnvironmentFactory envFactoryA = mock(EnvironmentFactory.class);
    when(envFactoryA.createEnvironment(eq(environmentA), any())).thenReturn(remoteEnvironment);
    EnvironmentFactory.Provider environmentProviderFactoryA = mock(EnvironmentFactory.Provider.class);
    when(environmentProviderFactoryA.createEnvironmentFactory(any(), any(), any(), any(), any(), any())).thenReturn(envFactoryA);
    when(environmentProviderFactoryA.getServerFactory()).thenReturn(serverFactory);
    Map<String, Provider> environmentFactoryProviderMap = ImmutableMap.of(environmentA.getUrn(), environmentProviderFactoryA);
    PortablePipelineOptions portableOptions = PipelineOptionsFactory.as(PortablePipelineOptions.class);
    portableOptions.setSdkWorkerParallelism(2);
    Struct pipelineOptions = PipelineOptionsTranslation.toProto(portableOptions);
    try (DefaultJobBundleFactory bundleFactory = new DefaultJobBundleFactory(JobInfo.create("testJob", "testJob", "token", pipelineOptions), environmentFactoryProviderMap, stageIdGenerator, serverInfo)) {
        bundleFactory.forStage(getExecutableStage(environmentA));
        verify(environmentProviderFactoryA, Mockito.times(1)).createEnvironmentFactory(any(), any(), any(), any(), any(), any());
        verify(envFactoryA, Mockito.times(1)).createEnvironment(eq(environmentA), any());
        bundleFactory.forStage(getExecutableStage(environmentA));
        verify(environmentProviderFactoryA, Mockito.times(2)).createEnvironmentFactory(any(), any(), any(), any(), any(), any());
        verify(envFactoryA, Mockito.times(2)).createEnvironment(eq(environmentA), any());
        // round robin, no new environment created
        bundleFactory.forStage(getExecutableStage(environmentA));
        verify(environmentProviderFactoryA, Mockito.times(2)).createEnvironmentFactory(any(), any(), any(), any(), any(), any());
        verify(envFactoryA, Mockito.times(2)).createEnvironment(eq(environmentA), any());
    }
    portableOptions.setSdkWorkerParallelism(0);
    pipelineOptions = PipelineOptionsTranslation.toProto(portableOptions);
    Mockito.reset(envFactoryA);
    when(envFactoryA.createEnvironment(eq(environmentA), any())).thenReturn(remoteEnvironment);
    int expectedParallelism = Math.max(1, Runtime.getRuntime().availableProcessors() - 1);
    try (DefaultJobBundleFactory bundleFactory = new DefaultJobBundleFactory(JobInfo.create("testJob", "testJob", "token", pipelineOptions), environmentFactoryProviderMap, stageIdGenerator, serverInfo)) {
        HashSet<StageBundleFactory> stageBundleFactorySet = new HashSet<>();
        // more factories than parallelism for round-robin
        int numStageBundleFactories = expectedParallelism + 5;
        for (int i = 0; i < numStageBundleFactories; i++) {
            stageBundleFactorySet.add(bundleFactory.forStage(getExecutableStage(environmentA)));
        }
        verify(envFactoryA, Mockito.times(expectedParallelism)).createEnvironment(eq(environmentA), any());
        Assert.assertEquals(numStageBundleFactories, stageBundleFactorySet.size());
    }
}

Example 7 with RemoteEnvironment

use of org.apache.beam.runners.fnexecution.environment.RemoteEnvironment in project beam by apache.

the class SingleEnvironmentInstanceJobBundleFactory method createBundleFactory.

private StageBundleFactory createBundleFactory(ExecutableStage stage) {
    RemoteEnvironment remoteEnv = environments.computeIfAbsent(stage.getEnvironment(), env -> {
        try {
            return environmentFactory.createEnvironment(env, idGenerator.getId());
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    });
    SdkHarnessClient sdkHarnessClient = SdkHarnessClient.usingFnApiClient(remoteEnv.getInstructionRequestHandler(), dataService.getService()).withIdGenerator(idGenerator);
    ExecutableProcessBundleDescriptor descriptor;
    try {
        descriptor = ProcessBundleDescriptors.fromExecutableStage(idGenerator.getId(), stage, dataService.getApiServiceDescriptor(), stateService.getApiServiceDescriptor());
    } catch (IOException e) {
        throw new RuntimeException(e);
    }
    SdkHarnessClient.BundleProcessor bundleProcessor = sdkHarnessClient.getProcessor(descriptor.getProcessBundleDescriptor(), descriptor.getRemoteInputDestinations(), stateService.getService());
    return new BundleProcessorStageBundleFactory(descriptor, bundleProcessor, sdkHarnessClient);
}

Example 8 with RemoteEnvironment

use of org.apache.beam.runners.fnexecution.environment.RemoteEnvironment in project beam by apache.

the class SingleEnvironmentInstanceJobBundleFactoryTest method closeShutsDownEnvironments.

@Test
public void closeShutsDownEnvironments() throws Exception {
    Pipeline p = Pipeline.create();
    ExperimentalOptions.addExperiment(p.getOptions().as(ExperimentalOptions.class), "beam_fn_api");
    p.apply("Create", Create.of(1, 2, 3));
    ExecutableStage stage = GreedyPipelineFuser.fuse(PipelineTranslation.toProto(p)).getFusedStages().stream().findFirst().get();
    RemoteEnvironment remoteEnv = mock(RemoteEnvironment.class);
    when(remoteEnv.getInstructionRequestHandler()).thenReturn(instructionRequestHandler);
    when(environmentFactory.createEnvironment(stage.getEnvironment(), GENERATED_ID)).thenReturn(remoteEnv);
    factory.forStage(stage);
    factory.close();
    verify(remoteEnv).close();
}