Search in sources :

Example 1 with BeamWorkerStatusGrpcService

use of org.apache.beam.runners.fnexecution.status.BeamWorkerStatusGrpcService in project beam by apache.

the class DataflowRunnerHarness method main.

/**
 * Fetches and processes work units from the Dataflow service.
 */
public static void main(String[] unusedArgs) throws Exception {
    RunnerApi.@Nullable Pipeline pipeline = DataflowWorkerHarnessHelper.getPipelineFromEnv();
    // This descriptor is used for all services except logging. They are isolated to keep
    // critical traffic protected from best effort traffic.
    ApiServiceDescriptor controlApiService = DataflowWorkerHarnessHelper.getControlDescriptor();
    ApiServiceDescriptor loggingApiService = DataflowWorkerHarnessHelper.getLoggingDescriptor();
    ApiServiceDescriptor statusApiService = DataflowWorkerHarnessHelper.getStatusDescriptor();
    LOG.info("{} started, using port {} for control, {} for logging.", DataflowRunnerHarness.class, controlApiService, loggingApiService);
    DataflowWorkerHarnessHelper.initializeLogging(DataflowRunnerHarness.class);
    DataflowWorkerHarnessOptions pipelineOptions = DataflowWorkerHarnessHelper.initializeGlobalStateAndPipelineOptions(DataflowRunnerHarness.class);
    DataflowWorkerHarnessHelper.configureLogging(pipelineOptions);
    // Initialized registered file systems.˜
    FileSystems.setDefaultPipelineOptions(pipelineOptions);
    DataflowPipelineDebugOptions dataflowOptions = pipelineOptions.as(DataflowPipelineDebugOptions.class);
    ServerFactory serverFactory;
    if (DataflowRunner.hasExperiment(dataflowOptions, "beam_fn_api_epoll_domain_socket")) {
        serverFactory = ServerFactory.createEpollDomainSocket();
    } else if (DataflowRunner.hasExperiment(dataflowOptions, "beam_fn_api_epoll")) {
        serverFactory = ServerFactory.createEpollSocket();
    } else {
        serverFactory = ServerFactory.createDefault();
    }
    ServerStreamObserverFactory streamObserverFactory = ServerStreamObserverFactory.fromOptions(pipelineOptions);
    Server servicesServer = null;
    Server loggingServer = null;
    Server statusServer = null;
    try (BeamFnLoggingService beamFnLoggingService = new BeamFnLoggingService(loggingApiService, DataflowWorkerLoggingInitializer.getSdkLoggingHandler()::publish, streamObserverFactory::from, GrpcContextHeaderAccessorProvider.getHeaderAccessor());
        BeamFnControlService beamFnControlService = new BeamFnControlService(controlApiService, streamObserverFactory::from, GrpcContextHeaderAccessorProvider.getHeaderAccessor());
        BeamFnDataGrpcService beamFnDataService = new BeamFnDataGrpcService(pipelineOptions, controlApiService, streamObserverFactory::from, GrpcContextHeaderAccessorProvider.getHeaderAccessor());
        BeamWorkerStatusGrpcService beamWorkerStatusGrpcService = statusApiService == null ? null : BeamWorkerStatusGrpcService.create(statusApiService, GrpcContextHeaderAccessorProvider.getHeaderAccessor());
        GrpcStateService beamFnStateService = GrpcStateService.create()) {
        servicesServer = serverFactory.create(ImmutableList.of(beamFnControlService, beamFnDataService, beamFnStateService), controlApiService);
        loggingServer = serverFactory.create(ImmutableList.of(beamFnLoggingService), loggingApiService);
        // gRPC server for obtaining SDK harness runtime status information.
        if (beamWorkerStatusGrpcService != null) {
            statusServer = serverFactory.create(ImmutableList.of(beamWorkerStatusGrpcService), statusApiService);
        }
        start(pipeline, pipelineOptions, beamFnControlService, beamFnDataService, controlApiService, beamFnStateService, beamWorkerStatusGrpcService);
        if (statusServer != null) {
            statusServer.shutdown();
        }
        servicesServer.shutdown();
        loggingServer.shutdown();
        // wait 30 secs for outstanding requests to finish.
        if (statusServer != null) {
            statusServer.awaitTermination(30, TimeUnit.SECONDS);
        }
        servicesServer.awaitTermination(30, TimeUnit.SECONDS);
        loggingServer.awaitTermination(30, TimeUnit.SECONDS);
    } finally {
        if (statusServer != null && !statusServer.isTerminated()) {
            statusServer.shutdownNow();
        }
        if (servicesServer != null && !servicesServer.isTerminated()) {
            servicesServer.shutdownNow();
        }
        if (loggingServer != null && !loggingServer.isTerminated()) {
            loggingServer.shutdownNow();
        }
    }
}
Also used : GrpcStateService(org.apache.beam.runners.fnexecution.state.GrpcStateService) ApiServiceDescriptor(org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor) Server(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server) DataflowWorkerHarnessOptions(org.apache.beam.runners.dataflow.options.DataflowWorkerHarnessOptions) ServerFactory(org.apache.beam.sdk.fn.server.ServerFactory) BeamFnControlService(org.apache.beam.runners.dataflow.worker.fn.BeamFnControlService) ServerStreamObserverFactory(org.apache.beam.runners.dataflow.worker.fn.stream.ServerStreamObserverFactory) BeamFnDataGrpcService(org.apache.beam.runners.dataflow.worker.fn.data.BeamFnDataGrpcService) RunnerApi(org.apache.beam.model.pipeline.v1.RunnerApi) BeamFnLoggingService(org.apache.beam.runners.dataflow.worker.fn.logging.BeamFnLoggingService) DataflowPipelineDebugOptions(org.apache.beam.runners.dataflow.options.DataflowPipelineDebugOptions) BeamWorkerStatusGrpcService(org.apache.beam.runners.fnexecution.status.BeamWorkerStatusGrpcService) Nullable(org.checkerframework.checker.nullness.qual.Nullable)

Aggregations

ApiServiceDescriptor (org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor)1 RunnerApi (org.apache.beam.model.pipeline.v1.RunnerApi)1 DataflowPipelineDebugOptions (org.apache.beam.runners.dataflow.options.DataflowPipelineDebugOptions)1 DataflowWorkerHarnessOptions (org.apache.beam.runners.dataflow.options.DataflowWorkerHarnessOptions)1 BeamFnControlService (org.apache.beam.runners.dataflow.worker.fn.BeamFnControlService)1 BeamFnDataGrpcService (org.apache.beam.runners.dataflow.worker.fn.data.BeamFnDataGrpcService)1 BeamFnLoggingService (org.apache.beam.runners.dataflow.worker.fn.logging.BeamFnLoggingService)1 ServerStreamObserverFactory (org.apache.beam.runners.dataflow.worker.fn.stream.ServerStreamObserverFactory)1 GrpcStateService (org.apache.beam.runners.fnexecution.state.GrpcStateService)1 BeamWorkerStatusGrpcService (org.apache.beam.runners.fnexecution.status.BeamWorkerStatusGrpcService)1 ServerFactory (org.apache.beam.sdk.fn.server.ServerFactory)1 Server (org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server)1 Nullable (org.checkerframework.checker.nullness.qual.Nullable)1