Search in sources :

Example 1 with GrpcStateService

use of org.apache.beam.runners.fnexecution.state.GrpcStateService in project beam by apache.

the class DefaultJobBundleFactoryTest method setUpMocks.

@Before
public void setUpMocks() throws Exception {
    MockitoAnnotations.initMocks(this);
    when(envFactory.createEnvironment(eq(environment), any())).thenReturn(remoteEnvironment);
    when(remoteEnvironment.getInstructionRequestHandler()).thenReturn(instructionHandler);
    when(instructionHandler.handle(any())).thenReturn(CompletableFuture.completedFuture(instructionResponse));
    when(dataServer.getApiServiceDescriptor()).thenReturn(ApiServiceDescriptor.getDefaultInstance());
    GrpcDataService dataService = mock(GrpcDataService.class);
    when(dataService.send(any(), any())).thenReturn(mock(CloseableFnDataReceiver.class));
    when(dataServer.getService()).thenReturn(dataService);
    when(stateServer.getApiServiceDescriptor()).thenReturn(ApiServiceDescriptor.getDefaultInstance());
    GrpcStateService stateService = mock(GrpcStateService.class);
    when(stateService.registerForProcessBundleInstructionId(any(), any())).thenReturn(mock(StateDelegator.Registration.class));
    when(stateServer.getService()).thenReturn(stateService);
    when(provisioningServer.getService()).thenReturn(provisionService);
    serverInfo = new AutoValue_DefaultJobBundleFactory_ServerInfo.Builder().setControlServer(controlServer).setLoggingServer(loggingServer).setRetrievalServer(retrievalServer).setProvisioningServer(provisioningServer).setDataServer(dataServer).setStateServer(stateServer).build();
}
Also used : CloseableFnDataReceiver(org.apache.beam.sdk.fn.data.CloseableFnDataReceiver) GrpcStateService(org.apache.beam.runners.fnexecution.state.GrpcStateService) GrpcDataService(org.apache.beam.runners.fnexecution.data.GrpcDataService) Before(org.junit.Before)

Example 2 with GrpcStateService

use of org.apache.beam.runners.fnexecution.state.GrpcStateService 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)

Example 3 with GrpcStateService

use of org.apache.beam.runners.fnexecution.state.GrpcStateService 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;
}
Also used : StaticGrpcProvisionService(org.apache.beam.runners.fnexecution.provisioning.StaticGrpcProvisionService) GrpcStateService(org.apache.beam.runners.fnexecution.state.GrpcStateService) GrpcDataService(org.apache.beam.runners.fnexecution.data.GrpcDataService) ArtifactRetrievalService(org.apache.beam.runners.fnexecution.artifact.ArtifactRetrievalService) GrpcLoggingService(org.apache.beam.runners.fnexecution.logging.GrpcLoggingService) PortablePipelineOptions(org.apache.beam.sdk.options.PortablePipelineOptions)

Aggregations

GrpcStateService (org.apache.beam.runners.fnexecution.state.GrpcStateService)3 GrpcDataService (org.apache.beam.runners.fnexecution.data.GrpcDataService)2 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 ArtifactRetrievalService (org.apache.beam.runners.fnexecution.artifact.ArtifactRetrievalService)1 GrpcLoggingService (org.apache.beam.runners.fnexecution.logging.GrpcLoggingService)1 StaticGrpcProvisionService (org.apache.beam.runners.fnexecution.provisioning.StaticGrpcProvisionService)1 BeamWorkerStatusGrpcService (org.apache.beam.runners.fnexecution.status.BeamWorkerStatusGrpcService)1 CloseableFnDataReceiver (org.apache.beam.sdk.fn.data.CloseableFnDataReceiver)1 ServerFactory (org.apache.beam.sdk.fn.server.ServerFactory)1 PortablePipelineOptions (org.apache.beam.sdk.options.PortablePipelineOptions)1 Server (org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server)1 Nullable (org.checkerframework.checker.nullness.qual.Nullable)1 Before (org.junit.Before)1