Examples with BeamFnStatusClient org.apache.beam.fn.harness.status.BeamFnStatusClient used on opensource projects

Example 1 with BeamFnStatusClient

use of org.apache.beam.fn.harness.status.BeamFnStatusClient 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();
    }
}