Search in sources :

Example 1 with HarnessMonitoringInfosInstructionHandler

use of org.apache.beam.fn.harness.control.HarnessMonitoringInfosInstructionHandler in project beam by apache.

the class FnHarness method main.

/**
 * Run a FnHarness with the given id and options that attaches to the specified logging and
 * control API service descriptors using the given channel factory and outbound observer factory.
 *
 * @param id Harness ID
 * @param options The options for this pipeline
 * @param runnerCapabilites
 * @param loggingApiServiceDescriptor
 * @param controlApiServiceDescriptor
 * @param statusApiServiceDescriptor
 * @param channelFactory
 * @param outboundObserverFactory
 * @param processWideCache
 * @throws Exception
 */
public static void main(String id, PipelineOptions options, Set<String> runnerCapabilites, Endpoints.ApiServiceDescriptor loggingApiServiceDescriptor, Endpoints.ApiServiceDescriptor controlApiServiceDescriptor, Endpoints.ApiServiceDescriptor statusApiServiceDescriptor, ManagedChannelFactory channelFactory, OutboundObserverFactory outboundObserverFactory, Cache<Object, Object> processWideCache) throws Exception {
    channelFactory = channelFactory.withInterceptors(ImmutableList.of(AddHarnessIdInterceptor.create(id)));
    IdGenerator idGenerator = IdGenerators.decrementingLongs();
    ShortIdMap metricsShortIds = new ShortIdMap();
    ExecutorService executorService = options.as(GcsOptions.class).getExecutorService();
    // intercepts logging and sends it to the logging service.
    try (BeamFnLoggingClient logging = new BeamFnLoggingClient(options, loggingApiServiceDescriptor, channelFactory::forDescriptor)) {
        LOG.info("Fn Harness started");
        // Register standard file systems.
        FileSystems.setDefaultPipelineOptions(options);
        EnumMap<BeamFnApi.InstructionRequest.RequestCase, ThrowingFunction<InstructionRequest, BeamFnApi.InstructionResponse.Builder>> handlers = new EnumMap<>(BeamFnApi.InstructionRequest.RequestCase.class);
        ManagedChannel channel = channelFactory.forDescriptor(controlApiServiceDescriptor);
        BeamFnControlGrpc.BeamFnControlStub controlStub = BeamFnControlGrpc.newStub(channel);
        BeamFnControlGrpc.BeamFnControlBlockingStub blockingControlStub = BeamFnControlGrpc.newBlockingStub(channel);
        BeamFnDataGrpcClient beamFnDataMultiplexer = new BeamFnDataGrpcClient(options, channelFactory::forDescriptor, outboundObserverFactory);
        BeamFnStateGrpcClientCache beamFnStateGrpcClientCache = new BeamFnStateGrpcClientCache(idGenerator, channelFactory, outboundObserverFactory);
        FinalizeBundleHandler finalizeBundleHandler = new FinalizeBundleHandler(options.as(GcsOptions.class).getExecutorService());
        Function<String, BeamFnApi.ProcessBundleDescriptor> getProcessBundleDescriptor = new Function<String, ProcessBundleDescriptor>() {

            private static final String PROCESS_BUNDLE_DESCRIPTORS = "ProcessBundleDescriptors";

            private final Cache<String, BeamFnApi.ProcessBundleDescriptor> cache = Caches.subCache(processWideCache, PROCESS_BUNDLE_DESCRIPTORS);

            @Override
            public BeamFnApi.ProcessBundleDescriptor apply(String id) {
                return cache.computeIfAbsent(id, this::loadDescriptor);
            }

            private BeamFnApi.ProcessBundleDescriptor loadDescriptor(String id) {
                return blockingControlStub.getProcessBundleDescriptor(BeamFnApi.GetProcessBundleDescriptorRequest.newBuilder().setProcessBundleDescriptorId(id).build());
            }
        };
        MetricsEnvironment.setProcessWideContainer(MetricsContainerImpl.createProcessWideContainer());
        ProcessBundleHandler processBundleHandler = new ProcessBundleHandler(options, runnerCapabilites, getProcessBundleDescriptor, beamFnDataMultiplexer, beamFnStateGrpcClientCache, finalizeBundleHandler, metricsShortIds, processWideCache);
        BeamFnStatusClient beamFnStatusClient = null;
        if (statusApiServiceDescriptor != null) {
            beamFnStatusClient = new BeamFnStatusClient(statusApiServiceDescriptor, channelFactory::forDescriptor, processBundleHandler.getBundleProcessorCache(), options, processWideCache);
        }
        // TODO(BEAM-9729): Remove once runners no longer send this instruction.
        handlers.put(BeamFnApi.InstructionRequest.RequestCase.REGISTER, request -> BeamFnApi.InstructionResponse.newBuilder().setRegister(BeamFnApi.RegisterResponse.getDefaultInstance()));
        handlers.put(BeamFnApi.InstructionRequest.RequestCase.FINALIZE_BUNDLE, finalizeBundleHandler::finalizeBundle);
        handlers.put(BeamFnApi.InstructionRequest.RequestCase.PROCESS_BUNDLE, processBundleHandler::processBundle);
        handlers.put(BeamFnApi.InstructionRequest.RequestCase.PROCESS_BUNDLE_PROGRESS, processBundleHandler::progress);
        handlers.put(BeamFnApi.InstructionRequest.RequestCase.PROCESS_BUNDLE_SPLIT, processBundleHandler::trySplit);
        handlers.put(InstructionRequest.RequestCase.MONITORING_INFOS, request -> BeamFnApi.InstructionResponse.newBuilder().setMonitoringInfos(BeamFnApi.MonitoringInfosMetadataResponse.newBuilder().putAllMonitoringInfo(StreamSupport.stream(request.getMonitoringInfos().getMonitoringInfoIdList().spliterator(), false).collect(Collectors.toMap(Function.identity(), metricsShortIds::get)))));
        HarnessMonitoringInfosInstructionHandler processWideHandler = new HarnessMonitoringInfosInstructionHandler(metricsShortIds);
        handlers.put(InstructionRequest.RequestCase.HARNESS_MONITORING_INFOS, processWideHandler::harnessMonitoringInfos);
        JvmInitializers.runBeforeProcessing(options);
        String samplingPeriodMills = ExperimentalOptions.getExperimentValue(options, ExperimentalOptions.STATE_SAMPLING_PERIOD_MILLIS);
        if (samplingPeriodMills != null) {
            ExecutionStateSampler.setSamplingPeriod(Integer.parseInt(samplingPeriodMills));
        }
        ExecutionStateSampler.instance().start();
        LOG.info("Entering instruction processing loop");
        // The control client immediately dispatches requests to an executor so we execute on the
        // direct executor. If we created separate channels for different stubs we could use
        // directExecutor() when building the channel.
        BeamFnControlClient control = new BeamFnControlClient(controlStub.withExecutor(MoreExecutors.directExecutor()), outboundObserverFactory, executorService, handlers);
        control.waitForTermination();
        if (beamFnStatusClient != null) {
            beamFnStatusClient.close();
        }
        processBundleHandler.shutdown();
    } finally {
        System.out.println("Shutting SDK harness down.");
        ExecutionStateSampler.instance().stop();
        executorService.shutdown();
    }
}
Also used : BeamFnControlClient(org.apache.beam.fn.harness.control.BeamFnControlClient) ProcessBundleDescriptor(org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleDescriptor) ProcessBundleHandler(org.apache.beam.fn.harness.control.ProcessBundleHandler) BeamFnStateGrpcClientCache(org.apache.beam.fn.harness.state.BeamFnStateGrpcClientCache) ThrowingFunction(org.apache.beam.sdk.function.ThrowingFunction) Function(java.util.function.Function) BeamFnStatusClient(org.apache.beam.fn.harness.status.BeamFnStatusClient) InstructionRequest(org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionRequest) ManagedChannel(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.ManagedChannel) GcsOptions(org.apache.beam.sdk.extensions.gcp.options.GcsOptions) EnumMap(java.util.EnumMap) BeamFnControlGrpc(org.apache.beam.model.fnexecution.v1.BeamFnControlGrpc) HarnessMonitoringInfosInstructionHandler(org.apache.beam.fn.harness.control.HarnessMonitoringInfosInstructionHandler) ThrowingFunction(org.apache.beam.sdk.function.ThrowingFunction) BeamFnDataGrpcClient(org.apache.beam.fn.harness.data.BeamFnDataGrpcClient) BeamFnApi(org.apache.beam.model.fnexecution.v1.BeamFnApi) IdGenerator(org.apache.beam.sdk.fn.IdGenerator) FinalizeBundleHandler(org.apache.beam.fn.harness.control.FinalizeBundleHandler) ShortIdMap(org.apache.beam.runners.core.metrics.ShortIdMap) ExecutorService(java.util.concurrent.ExecutorService) BeamFnLoggingClient(org.apache.beam.fn.harness.logging.BeamFnLoggingClient) BeamFnStateGrpcClientCache(org.apache.beam.fn.harness.state.BeamFnStateGrpcClientCache)

Aggregations

EnumMap (java.util.EnumMap)1 ExecutorService (java.util.concurrent.ExecutorService)1 Function (java.util.function.Function)1 BeamFnControlClient (org.apache.beam.fn.harness.control.BeamFnControlClient)1 FinalizeBundleHandler (org.apache.beam.fn.harness.control.FinalizeBundleHandler)1 HarnessMonitoringInfosInstructionHandler (org.apache.beam.fn.harness.control.HarnessMonitoringInfosInstructionHandler)1 ProcessBundleHandler (org.apache.beam.fn.harness.control.ProcessBundleHandler)1 BeamFnDataGrpcClient (org.apache.beam.fn.harness.data.BeamFnDataGrpcClient)1 BeamFnLoggingClient (org.apache.beam.fn.harness.logging.BeamFnLoggingClient)1 BeamFnStateGrpcClientCache (org.apache.beam.fn.harness.state.BeamFnStateGrpcClientCache)1 BeamFnStatusClient (org.apache.beam.fn.harness.status.BeamFnStatusClient)1 BeamFnApi (org.apache.beam.model.fnexecution.v1.BeamFnApi)1 InstructionRequest (org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionRequest)1 ProcessBundleDescriptor (org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleDescriptor)1 BeamFnControlGrpc (org.apache.beam.model.fnexecution.v1.BeamFnControlGrpc)1 ShortIdMap (org.apache.beam.runners.core.metrics.ShortIdMap)1 GcsOptions (org.apache.beam.sdk.extensions.gcp.options.GcsOptions)1 IdGenerator (org.apache.beam.sdk.fn.IdGenerator)1 ThrowingFunction (org.apache.beam.sdk.function.ThrowingFunction)1 ManagedChannel (org.apache.beam.vendor.grpc.v1p43p2.io.grpc.ManagedChannel)1