Search in sources :

Example 16 with ApiServiceDescriptor

use of org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor in project beam by apache.

the class GrpcFnServer method allocatePortAndCreateFor.

/**
 * Create a {@link GrpcFnServer} for the provided {@link FnService} running on an arbitrary port.
 */
public static <ServiceT extends FnService> GrpcFnServer<ServiceT> allocatePortAndCreateFor(ServiceT service, ServerFactory factory) throws IOException {
    ApiServiceDescriptor.Builder apiServiceDescriptor = ApiServiceDescriptor.newBuilder();
    Server server = factory.allocateAddressAndCreate(ImmutableList.of(service), apiServiceDescriptor);
    return new GrpcFnServer<>(server, service, apiServiceDescriptor.build());
}
Also used : ApiServiceDescriptor(org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor) Server(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server)

Example 17 with ApiServiceDescriptor

use of org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor 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 18 with ApiServiceDescriptor

use of org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor in project beam by apache.

the class PortableRunner method run.

@Override
public PipelineResult run(Pipeline pipeline) {
    Runnable cleanup;
    if (Environments.ENVIRONMENT_LOOPBACK.equals(options.as(PortablePipelineOptions.class).getDefaultEnvironmentType())) {
        GrpcFnServer<ExternalWorkerService> workerService;
        try {
            workerService = new ExternalWorkerService(options).start();
        } catch (Exception exn) {
            throw new RuntimeException("Failed to start GrpcFnServer for ExternalWorkerService", exn);
        }
        LOG.info("Starting worker service at {}", workerService.getApiServiceDescriptor().getUrl());
        options.as(PortablePipelineOptions.class).setDefaultEnvironmentConfig(workerService.getApiServiceDescriptor().getUrl());
        cleanup = () -> {
            try {
                LOG.warn("closing worker service {}", workerService);
                workerService.close();
            } catch (Exception exn) {
                throw new RuntimeException(exn);
            }
        };
    } else {
        cleanup = null;
    }
    ImmutableList.Builder<String> filesToStageBuilder = ImmutableList.builder();
    List<String> stagingFiles = options.as(PortablePipelineOptions.class).getFilesToStage();
    if (stagingFiles == null) {
        List<String> classpathResources = detectClassPathResourcesToStage(Environments.class.getClassLoader(), options);
        if (classpathResources.isEmpty()) {
            throw new IllegalArgumentException("No classpath elements found.");
        }
        LOG.debug("PortablePipelineOptions.filesToStage was not specified. " + "Defaulting to files from the classpath: {}", classpathResources.size());
        filesToStageBuilder.addAll(classpathResources);
    } else {
        filesToStageBuilder.addAll(stagingFiles);
    }
    // TODO(heejong): remove jar_packages experimental flag when cross-language dependency
    // management is implemented for all runners.
    List<String> experiments = options.as(ExperimentalOptions.class).getExperiments();
    if (experiments != null) {
        Optional<String> jarPackages = experiments.stream().filter((String flag) -> flag.startsWith("jar_packages=")).findFirst();
        jarPackages.ifPresent(s -> filesToStageBuilder.addAll(Arrays.asList(s.replaceFirst("jar_packages=", "").split(","))));
    }
    options.as(PortablePipelineOptions.class).setFilesToStage(filesToStageBuilder.build());
    RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline, SdkComponents.create(options));
    pipelineProto = DefaultArtifactResolver.INSTANCE.resolveArtifacts(pipelineProto);
    PrepareJobRequest prepareJobRequest = PrepareJobRequest.newBuilder().setJobName(options.getJobName()).setPipeline(pipelineProto).setPipelineOptions(PipelineOptionsTranslation.toProto(options)).build();
    LOG.info("Using job server endpoint: {}", endpoint);
    ManagedChannel jobServiceChannel = channelFactory.forDescriptor(ApiServiceDescriptor.newBuilder().setUrl(endpoint).build());
    JobServiceBlockingStub jobService = JobServiceGrpc.newBlockingStub(jobServiceChannel);
    try (CloseableResource<JobServiceBlockingStub> wrappedJobService = CloseableResource.of(jobService, unused -> jobServiceChannel.shutdown())) {
        final int jobServerTimeout = options.as(PortablePipelineOptions.class).getJobServerTimeout();
        PrepareJobResponse prepareJobResponse = jobService.withDeadlineAfter(jobServerTimeout, TimeUnit.SECONDS).withWaitForReady().prepare(prepareJobRequest);
        LOG.info("PrepareJobResponse: {}", prepareJobResponse);
        ApiServiceDescriptor artifactStagingEndpoint = prepareJobResponse.getArtifactStagingEndpoint();
        String stagingSessionToken = prepareJobResponse.getStagingSessionToken();
        try (CloseableResource<ManagedChannel> artifactChannel = CloseableResource.of(channelFactory.forDescriptor(artifactStagingEndpoint), ManagedChannel::shutdown)) {
            ArtifactStagingService.offer(new ArtifactRetrievalService(), ArtifactStagingServiceGrpc.newStub(artifactChannel.get()), stagingSessionToken);
        } catch (CloseableResource.CloseException e) {
            LOG.warn("Error closing artifact staging channel", e);
        // CloseExceptions should only be thrown while closing the channel.
        } catch (Exception e) {
            throw new RuntimeException("Error staging files.", e);
        }
        RunJobRequest runJobRequest = RunJobRequest.newBuilder().setPreparationId(prepareJobResponse.getPreparationId()).build();
        // Run the job and wait for a result, we don't set a timeout here because
        // it may take a long time for a job to complete and streaming
        // jobs never return a response.
        RunJobResponse runJobResponse = jobService.run(runJobRequest);
        LOG.info("RunJobResponse: {}", runJobResponse);
        ByteString jobId = runJobResponse.getJobIdBytes();
        return new JobServicePipelineResult(jobId, jobServerTimeout, wrappedJobService.transfer(), cleanup);
    } catch (CloseException e) {
        throw new RuntimeException(e);
    }
}
Also used : JobServiceBlockingStub(org.apache.beam.model.jobmanagement.v1.JobServiceGrpc.JobServiceBlockingStub) PrepareJobResponse(org.apache.beam.model.jobmanagement.v1.JobApi.PrepareJobResponse) ApiServiceDescriptor(org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor) ImmutableList(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) CloseException(org.apache.beam.runners.portability.CloseableResource.CloseException) ExperimentalOptions(org.apache.beam.sdk.options.ExperimentalOptions) ArtifactRetrievalService(org.apache.beam.runners.fnexecution.artifact.ArtifactRetrievalService) RunJobResponse(org.apache.beam.model.jobmanagement.v1.JobApi.RunJobResponse) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) RunnerApi(org.apache.beam.model.pipeline.v1.RunnerApi) RunJobRequest(org.apache.beam.model.jobmanagement.v1.JobApi.RunJobRequest) ExternalWorkerService(org.apache.beam.fn.harness.ExternalWorkerService) PrepareJobRequest(org.apache.beam.model.jobmanagement.v1.JobApi.PrepareJobRequest) ManagedChannel(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.ManagedChannel) CloseException(org.apache.beam.runners.portability.CloseableResource.CloseException) CloseException(org.apache.beam.runners.portability.CloseableResource.CloseException) Environments(org.apache.beam.runners.core.construction.Environments) PortablePipelineOptions(org.apache.beam.sdk.options.PortablePipelineOptions)

Example 19 with ApiServiceDescriptor

use of org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor in project beam by apache.

the class BeamFnLoggingClientTest method testWhenServerFailsThatClientIsAbleToCleanup.

@Test
public void testWhenServerFailsThatClientIsAbleToCleanup() throws Exception {
    BeamFnLoggingMDC.setInstructionId("instruction-1");
    Collection<BeamFnApi.LogEntry> values = new ConcurrentLinkedQueue<>();
    AtomicReference<StreamObserver<BeamFnApi.LogControl>> outboundServerObserver = new AtomicReference<>();
    CallStreamObserver<BeamFnApi.LogEntry.List> inboundServerObserver = TestStreams.withOnNext((BeamFnApi.LogEntry.List logEntries) -> values.addAll(logEntries.getLogEntriesList())).build();
    Endpoints.ApiServiceDescriptor apiServiceDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl(this.getClass().getName() + "-" + UUID.randomUUID().toString()).build();
    Server server = InProcessServerBuilder.forName(apiServiceDescriptor.getUrl()).addService(new BeamFnLoggingGrpc.BeamFnLoggingImplBase() {

        @Override
        public StreamObserver<BeamFnApi.LogEntry.List> logging(StreamObserver<BeamFnApi.LogControl> outboundObserver) {
            outboundServerObserver.set(outboundObserver);
            outboundObserver.onError(Status.INTERNAL.withDescription("TEST ERROR").asException());
            return inboundServerObserver;
        }
    }).build();
    server.start();
    ManagedChannel channel = InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build();
    // Keep a strong reference to the loggers. Otherwise the call to client.close()
    // removes the only reference and the logger may get GC'd before the assertions (BEAM-4136).
    Logger rootLogger = null;
    Logger configuredLogger = null;
    try {
        BeamFnLoggingClient client = new BeamFnLoggingClient(PipelineOptionsFactory.fromArgs(new String[] { "--defaultSdkHarnessLogLevel=OFF", "--sdkHarnessLogLevelOverrides={\"ConfiguredLogger\": \"DEBUG\"}" }).create(), apiServiceDescriptor, (Endpoints.ApiServiceDescriptor descriptor) -> channel);
        rootLogger = LogManager.getLogManager().getLogger("");
        configuredLogger = LogManager.getLogManager().getLogger("ConfiguredLogger");
        thrown.expectMessage("TEST ERROR");
        client.close();
    } finally {
        assertNotNull("rootLogger should be initialized before exception", rootLogger);
        assertNotNull("configuredLogger should be initialized before exception", rootLogger);
        // Verify that after close, log levels are reset.
        assertEquals(Level.INFO, rootLogger.getLevel());
        assertNull(configuredLogger.getLevel());
        assertTrue(channel.isShutdown());
        server.shutdownNow();
    }
}
Also used : CallStreamObserver(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.CallStreamObserver) StreamObserver(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver) Server(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server) BeamFnApi(org.apache.beam.model.fnexecution.v1.BeamFnApi) AtomicReference(java.util.concurrent.atomic.AtomicReference) Logger(java.util.logging.Logger) Endpoints(org.apache.beam.model.pipeline.v1.Endpoints) ManagedChannel(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.ManagedChannel) ConcurrentLinkedQueue(java.util.concurrent.ConcurrentLinkedQueue) Test(org.junit.Test)

Example 20 with ApiServiceDescriptor

use of org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor in project beam by apache.

the class PTransformRunnerFactoryTestContext method addOutgoingDataEndpoint.

@Override
public <T> FnDataReceiver<T> addOutgoingDataEndpoint(ApiServiceDescriptor apiServiceDescriptor, Coder<T> coder) {
    BeamFnDataOutboundAggregator aggregator = getOutboundAggregators().get(apiServiceDescriptor);
    FnDataReceiver<T> receiver = aggregator.registerOutputDataLocation(getPTransformId(), coder);
    getOutgoingDataEndpoints().computeIfAbsent(apiServiceDescriptor, (unused) -> new ArrayList<>()).add(DataEndpoint.create(getPTransformId(), coder, receiver));
    return receiver;
}
Also used : BeamFnDataOutboundAggregator(org.apache.beam.sdk.fn.data.BeamFnDataOutboundAggregator) BeamFnDataOutboundAggregator(org.apache.beam.sdk.fn.data.BeamFnDataOutboundAggregator) Elements(org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements) WindowedValue(org.apache.beam.sdk.util.WindowedValue) BeamFnStateClient(org.apache.beam.fn.harness.state.BeamFnStateClient) BundleFinalizer(org.apache.beam.sdk.transforms.DoFn.BundleFinalizer) Coder(org.apache.beam.sdk.coders.Coder) HashMap(java.util.HashMap) CompletableFuture(java.util.concurrent.CompletableFuture) PipelineOptionsFactory(org.apache.beam.sdk.options.PipelineOptionsFactory) Supplier(java.util.function.Supplier) ArrayList(java.util.ArrayList) HashSet(java.util.HashSet) Map(java.util.Map) DelayedBundleApplication(org.apache.beam.model.fnexecution.v1.BeamFnApi.DelayedBundleApplication) BeamFnDataClient(org.apache.beam.fn.harness.data.BeamFnDataClient) BundleApplication(org.apache.beam.model.fnexecution.v1.BeamFnApi.BundleApplication) NoSuchElementException(java.util.NoSuchElementException) MetricsApi(org.apache.beam.model.pipeline.v1.MetricsApi) ApiServiceDescriptor(org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor) PipelineOptions(org.apache.beam.sdk.options.PipelineOptions) RunnerApi(org.apache.beam.model.pipeline.v1.RunnerApi) DoFn(org.apache.beam.sdk.transforms.DoFn) CloseableFnDataReceiver(org.apache.beam.sdk.fn.data.CloseableFnDataReceiver) TimerEndpoint(org.apache.beam.sdk.fn.data.TimerEndpoint) FnDataReceiver(org.apache.beam.sdk.fn.data.FnDataReceiver) Set(java.util.Set) BundleSplitListener(org.apache.beam.fn.harness.control.BundleSplitListener) BeamFnApi(org.apache.beam.model.fnexecution.v1.BeamFnApi) ProgressRequestCallback(org.apache.beam.fn.harness.PTransformRunnerFactory.ProgressRequestCallback) DataEndpoint(org.apache.beam.sdk.fn.data.DataEndpoint) List(java.util.List) StateRequest(org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest) Timer(org.apache.beam.runners.core.construction.Timer) Instant(org.joda.time.Instant) AutoValue(com.google.auto.value.AutoValue) Collections(java.util.Collections) StateResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.StateResponse) ThrowingRunnable(org.apache.beam.sdk.function.ThrowingRunnable) ArrayList(java.util.ArrayList)

Aggregations

Test (org.junit.Test)18 Endpoints (org.apache.beam.model.pipeline.v1.Endpoints)16 Server (org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server)12 BeamFnApi (org.apache.beam.model.fnexecution.v1.BeamFnApi)11 ApiServiceDescriptor (org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor)11 ManagedChannel (org.apache.beam.vendor.grpc.v1p43p2.io.grpc.ManagedChannel)10 StreamObserver (org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver)10 CallStreamObserver (org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.CallStreamObserver)9 ConcurrentLinkedQueue (java.util.concurrent.ConcurrentLinkedQueue)6 AtomicReference (java.util.concurrent.atomic.AtomicReference)6 WindowedValue (org.apache.beam.sdk.util.WindowedValue)6 Coder (org.apache.beam.sdk.coders.Coder)5 PTransform (org.apache.beam.model.pipeline.v1.RunnerApi.PTransform)4 ArrayList (java.util.ArrayList)3 LinkedHashMap (java.util.LinkedHashMap)3 Map (java.util.Map)3 CountDownLatch (java.util.concurrent.CountDownLatch)3 Logger (java.util.logging.Logger)3 RunnerApi (org.apache.beam.model.pipeline.v1.RunnerApi)3 ByteStringCoder (org.apache.beam.runners.fnexecution.wire.ByteStringCoder)3