Examples with PortablePipelineOptions org.apache.beam.sdk.options.PortablePipelineOptions used on opensource projects

Example 16 with PortablePipelineOptions

use of org.apache.beam.sdk.options.PortablePipelineOptions 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 17 with PortablePipelineOptions

use of org.apache.beam.sdk.options.PortablePipelineOptions in project beam by apache.

the class DefaultJobBundleFactoryTest method rejectsStateCachingWithLoadBalancing.

@Test
public void rejectsStateCachingWithLoadBalancing() throws Exception {
    PortablePipelineOptions portableOptions = PipelineOptionsFactory.as(PortablePipelineOptions.class);
    portableOptions.setLoadBalanceBundles(true);
    ExperimentalOptions options = portableOptions.as(ExperimentalOptions.class);
    ExperimentalOptions.addExperiment(options, "state_cache_size=1");
    Struct pipelineOptions = PipelineOptionsTranslation.toProto(options);
    Exception e = Assert.assertThrows(IllegalArgumentException.class, () -> new DefaultJobBundleFactory(JobInfo.create("testJob", "testJob", "token", pipelineOptions), envFactoryProviderMap, stageIdGenerator, serverInfo).close());
    assertThat(e.getMessage(), containsString("state_cache_size"));
}

Example 18 with PortablePipelineOptions

use of org.apache.beam.sdk.options.PortablePipelineOptions in project beam by apache.

the class TestPortableRunner method run.

@Override
public PipelineResult run(Pipeline pipeline) {
    TestPortablePipelineOptions testPortablePipelineOptions = options.as(TestPortablePipelineOptions.class);
    String jobServerHostPort;
    JobServerDriver jobServerDriver;
    Class<JobServerDriver> jobServerDriverClass = testPortablePipelineOptions.getJobServerDriver();
    String[] parameters = testPortablePipelineOptions.getJobServerConfig();
    try {
        jobServerDriver = InstanceBuilder.ofType(jobServerDriverClass).fromFactoryMethod("fromParams").withArg(String[].class, parameters).build();
        jobServerHostPort = jobServerDriver.start();
    } catch (IOException e) {
        throw new RuntimeException("Failed to start job server", e);
    }
    try {
        PortablePipelineOptions portableOptions = options.as(PortablePipelineOptions.class);
        portableOptions.setRunner(PortableRunner.class);
        portableOptions.setJobEndpoint(jobServerHostPort);
        PortableRunner runner = PortableRunner.fromOptions(portableOptions);
        PipelineResult result = runner.run(pipeline);
        assertThat("Pipeline did not succeed.", result.waitUntilFinish(), Matchers.is(State.DONE));
        return result;
    } finally {
        jobServerDriver.stop();
    }
}

Example 19 with PortablePipelineOptions

use of org.apache.beam.sdk.options.PortablePipelineOptions in project beam by apache.

the class PortableRunnerTest method createPipelineOptions.

private static PipelineOptions createPipelineOptions() {
    PortablePipelineOptions options = PipelineOptionsFactory.create().as(PortablePipelineOptions.class);
    options.setJobEndpoint(ENDPOINT_URL);
    options.setRunner(PortableRunner.class);
    return options;
}

Example 20 with PortablePipelineOptions

use of org.apache.beam.sdk.options.PortablePipelineOptions in project beam by apache.

the class PortableRunner method create.

@VisibleForTesting
static PortableRunner create(PipelineOptions options, ManagedChannelFactory channelFactory) {
    PortablePipelineOptions portableOptions = PipelineOptionsValidator.validate(PortablePipelineOptions.class, options);
    String endpoint = portableOptions.getJobEndpoint();
    return new PortableRunner(options, endpoint, channelFactory);
}