use of org.apache.beam.runners.core.metrics.ShortIdMap in project beam by apache.
the class ProcessBundleHandlerTest method testTrySplitBeforeBundleDoesNotFail.
@Test
public void testTrySplitBeforeBundleDoesNotFail() {
ProcessBundleHandler handler = new ProcessBundleHandler(PipelineOptionsFactory.create(), Collections.emptySet(), null, beamFnDataClient, null, /* beamFnStateClient */
null, /* finalizeBundleHandler */
new ShortIdMap(), ImmutableMap.of(), Caches.noop(), new BundleProcessorCache());
BeamFnApi.InstructionResponse response = handler.trySplit(BeamFnApi.InstructionRequest.newBuilder().setInstructionId("999L").setProcessBundleSplit(BeamFnApi.ProcessBundleSplitRequest.newBuilder().setInstructionId("unknown-id")).build()).build();
assertNotNull(response.getProcessBundleSplit());
assertEquals(0, response.getProcessBundleSplit().getChannelSplitsCount());
}
use of org.apache.beam.runners.core.metrics.ShortIdMap 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();
}
}
Aggregations