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

Example 1 with PortablePipelineOptions

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

the class DefaultJobBundleFactoryTest method loadBalancesBundles.

@Test
public void loadBalancesBundles() throws Exception {
    PortablePipelineOptions portableOptions = PipelineOptionsFactory.as(PortablePipelineOptions.class);
    portableOptions.setSdkWorkerParallelism(2);
    portableOptions.setLoadBalanceBundles(true);
    Struct pipelineOptions = PipelineOptionsTranslation.toProto(portableOptions);
    try (DefaultJobBundleFactory bundleFactory = new DefaultJobBundleFactory(JobInfo.create("testJob", "testJob", "token", pipelineOptions), envFactoryProviderMap, stageIdGenerator, serverInfo)) {
        OutputReceiverFactory orf = mock(OutputReceiverFactory.class);
        StateRequestHandler srh = mock(StateRequestHandler.class);
        when(srh.getCacheTokens()).thenReturn(Collections.emptyList());
        StageBundleFactory sbf = bundleFactory.forStage(getExecutableStage(environment));
        RemoteBundle b1 = sbf.getBundle(orf, srh, BundleProgressHandler.ignored());
        verify(envFactory, Mockito.times(1)).createEnvironment(eq(environment), any());
        final RemoteBundle b2 = sbf.getBundle(orf, srh, BundleProgressHandler.ignored());
        verify(envFactory, Mockito.times(2)).createEnvironment(eq(environment), any());
        AtomicBoolean b2Closing = new AtomicBoolean(false);
        ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
        ScheduledFuture<Optional<Exception>> closingFuture = executor.schedule(() -> {
            try {
                b2Closing.compareAndSet(false, true);
                b2.close();
                return Optional.empty();
            } catch (Exception e) {
                return Optional.of(e);
            }
        }, 100, TimeUnit.MILLISECONDS);
        assertThat(b2Closing.get(), equalTo(false));
        // This call should block until closingFuture has finished closing b2 (100ms)
        RemoteBundle b3 = sbf.getBundle(orf, srh, BundleProgressHandler.ignored());
        // ensure the previous call waited for close
        assertThat(b2Closing.get(), equalTo(true));
        // Join closingFuture and check if an exception occurred
        Optional<Exception> closingException = closingFuture.get();
        if (closingException.isPresent()) {
            throw new AssertionError("Exception occurred while closing b2", closingException.get());
        }
        verify(envFactory, Mockito.times(2)).createEnvironment(eq(environment), any());
        b3.close();
        b1.close();
        executor.shutdown();
    }
}

Example 2 with PortablePipelineOptions

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

the class DefaultJobBundleFactory method getMaxEnvironmentClients.

private static int getMaxEnvironmentClients(JobInfo jobInfo) {
    PortablePipelineOptions pipelineOptions = PipelineOptionsTranslation.fromProto(jobInfo.pipelineOptions()).as(PortablePipelineOptions.class);
    int maxEnvironments = MoreObjects.firstNonNull(pipelineOptions.getSdkWorkerParallelism(), 1);
    Preconditions.checkArgument(maxEnvironments >= 0, "sdk_worker_parallelism must be >= 0");
    if (maxEnvironments == 0) {
        // if this is 0, use the auto behavior of num_cores - 1 so that we leave some resources
        // available for the java process
        maxEnvironments = Math.max(Runtime.getRuntime().availableProcessors() - 1, 1);
    }
    return maxEnvironments;
}

Example 3 with PortablePipelineOptions

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

the class DefaultJobBundleFactory method createServerInfo.

private ServerInfo createServerInfo(JobInfo jobInfo, ServerFactory serverFactory) throws IOException {
    Preconditions.checkNotNull(serverFactory, "serverFactory can not be null");
    PortablePipelineOptions portableOptions = PipelineOptionsTranslation.fromProto(jobInfo.pipelineOptions()).as(PortablePipelineOptions.class);
    GrpcFnServer<FnApiControlClientPoolService> controlServer = GrpcFnServer.allocatePortAndCreateFor(FnApiControlClientPoolService.offeringClientsToPool(clientPool.getSink(), GrpcContextHeaderAccessorProvider.getHeaderAccessor()), serverFactory);
    GrpcFnServer<GrpcLoggingService> loggingServer = GrpcFnServer.allocatePortAndCreateFor(GrpcLoggingService.forWriter(Slf4jLogWriter.getDefault()), serverFactory);
    GrpcFnServer<ArtifactRetrievalService> retrievalServer = GrpcFnServer.allocatePortAndCreateFor(new ArtifactRetrievalService(), serverFactory);
    ProvisionApi.ProvisionInfo.Builder provisionInfo = jobInfo.toProvisionInfo().toBuilder();
    provisionInfo.setLoggingEndpoint(loggingServer.getApiServiceDescriptor());
    provisionInfo.setArtifactEndpoint(retrievalServer.getApiServiceDescriptor());
    provisionInfo.setControlEndpoint(controlServer.getApiServiceDescriptor());
    GrpcFnServer<StaticGrpcProvisionService> provisioningServer = GrpcFnServer.allocatePortAndCreateFor(StaticGrpcProvisionService.create(provisionInfo.build(), GrpcContextHeaderAccessorProvider.getHeaderAccessor()), serverFactory);
    GrpcFnServer<GrpcDataService> dataServer = GrpcFnServer.allocatePortAndCreateFor(GrpcDataService.create(portableOptions, executor, OutboundObserverFactory.serverDirect()), serverFactory);
    GrpcFnServer<GrpcStateService> stateServer = GrpcFnServer.allocatePortAndCreateFor(GrpcStateService.create(), serverFactory);
    ServerInfo serverInfo = new AutoValue_DefaultJobBundleFactory_ServerInfo.Builder().setControlServer(controlServer).setLoggingServer(loggingServer).setRetrievalServer(retrievalServer).setProvisioningServer(provisioningServer).setDataServer(dataServer).setStateServer(stateServer).build();
    return serverInfo;
}

Example 4 with PortablePipelineOptions

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

the class FlinkJobInvoker method invokeWithExecutor.

@Override
protected JobInvocation invokeWithExecutor(RunnerApi.Pipeline pipeline, Struct options, @Nullable String retrievalToken, ListeningExecutorService executorService) {
    // TODO: How to make Java/Python agree on names of keys and their values?
    LOG.trace("Parsing pipeline options");
    FlinkPipelineOptions flinkOptions = PipelineOptionsTranslation.fromProto(options).as(FlinkPipelineOptions.class);
    String invocationId = String.format("%s_%s", flinkOptions.getJobName(), UUID.randomUUID().toString());
    if (FlinkPipelineOptions.AUTO.equals(flinkOptions.getFlinkMaster())) {
        flinkOptions.setFlinkMaster(serverConfig.getFlinkMaster());
    }
    PortablePipelineOptions portableOptions = flinkOptions.as(PortablePipelineOptions.class);
    PortablePipelineRunner pipelineRunner;
    if (Strings.isNullOrEmpty(portableOptions.getOutputExecutablePath())) {
        pipelineRunner = new FlinkPipelineRunner(flinkOptions, serverConfig.getFlinkConfDir(), detectClassPathResourcesToStage(FlinkJobInvoker.class.getClassLoader(), flinkOptions));
    } else {
        pipelineRunner = new PortablePipelineJarCreator(FlinkPipelineRunner.class);
    }
    flinkOptions.setRunner(null);
    LOG.info("Invoking job {} with pipeline runner {}", invocationId, pipelineRunner);
    return createJobInvocation(invocationId, retrievalToken, executorService, pipeline, flinkOptions, pipelineRunner);
}

Example 5 with PortablePipelineOptions

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

the class EnvironmentsTest method createEnvironmentDockerFromEnvironmentOptions.

@Test
public void createEnvironmentDockerFromEnvironmentOptions() {
    PortablePipelineOptions options = PipelineOptionsFactory.as(PortablePipelineOptions.class);
    options.setDefaultEnvironmentType(Environments.ENVIRONMENT_DOCKER);
    options.setEnvironmentOptions(ImmutableList.of("docker_container_image=java"));
    assertThat(Environments.createOrGetDefaultEnvironment(options), is(Environment.newBuilder().setUrn(BeamUrns.getUrn(StandardEnvironments.Environments.DOCKER)).setPayload(DockerPayload.newBuilder().setContainerImage("java").build().toByteString()).addAllCapabilities(Environments.getJavaCapabilities()).build()));
}