def test_shared_shutdown_does_nothing(self): thread_pool_executor.shared_unbounded_instance().shutdown() futures = [] with thread_pool_executor.shared_unbounded_instance() as executor: for _ in range(0, 5): futures.append(executor.submit(self.append_and_sleep, 0.01)) for future in futures: future.result(timeout=10) with self._lock: self.assertEqual(5, len(self._worker_idents))
def start( cls, use_process=False, port=0, state_cache_size=0, data_buffer_time_limit_ms=-1, container_executable=None # type: Optional[str] ): # type: (...) -> Tuple[str, grpc.Server] options = [("grpc.http2.max_pings_without_data", 0), ("grpc.http2.max_ping_strikes", 0)] worker_server = grpc.server( thread_pool_executor.shared_unbounded_instance(), options=options) worker_address = 'localhost:%s' % worker_server.add_insecure_port( '[::]:%s' % port) worker_pool = cls(use_process=use_process, container_executable=container_executable, state_cache_size=state_cache_size, data_buffer_time_limit_ms=data_buffer_time_limit_ms) beam_fn_api_pb2_grpc.add_BeamFnExternalWorkerPoolServicer_to_server( worker_pool, worker_server) worker_server.start() _LOGGER.info('Listening for workers at %s', worker_address) # Register to kill the subprocesses on exit. def kill_worker_processes(): for worker_process in worker_pool._worker_processes.values(): worker_process.kill() atexit.register(kill_worker_processes) return worker_address, worker_server
def main(argv): parser = argparse.ArgumentParser() parser.add_argument('-p', '--port', type=int, help='port on which to serve the job api') parser.add_argument('--fully_qualified_name_glob', default=None) known_args, pipeline_args = parser.parse_known_args(argv) pipeline_options = PipelineOptions( pipeline_args + ["--experiments=beam_fn_api", "--sdk_location=container"]) with fully_qualified_named_transform.FullyQualifiedNamedTransform.with_filter( known_args.fully_qualified_name_glob): server = grpc.server(thread_pool_executor.shared_unbounded_instance()) beam_expansion_api_pb2_grpc.add_ExpansionServiceServicer_to_server( expansion_service.ExpansionServiceServicer(pipeline_options), server) beam_artifact_api_pb2_grpc.add_ArtifactRetrievalServiceServicer_to_server( artifact_service.ArtifactRetrievalService( artifact_service.BeamFilesystemHandler(None).file_reader), server) server.add_insecure_port('localhost:{}'.format(known_args.port)) server.start() _LOGGER.info('Listening for expansion requests at %d', known_args.port) def cleanup(unused_signum, unused_frame): _LOGGER.info('Shutting down expansion service.') server.stop(None) signal.signal(signal.SIGTERM, cleanup) signal.signal(signal.SIGINT, cleanup) # blocking main thread forever. signal.pause()
def main(unused_argv): PyPIArtifactRegistry.register_artifact('beautifulsoup4', '>=4.9,<5.0') parser = argparse.ArgumentParser() parser.add_argument('-p', '--port', type=int, help='port on which to serve the job api') parser.add_argument('--fully_qualified_name_glob', default=None) options = parser.parse_args() global server with fully_qualified_named_transform.FullyQualifiedNamedTransform.with_filter( options.fully_qualified_name_glob): server = grpc.server(thread_pool_executor.shared_unbounded_instance()) beam_expansion_api_pb2_grpc.add_ExpansionServiceServicer_to_server( expansion_service.ExpansionServiceServicer( PipelineOptions([ "--experiments", "beam_fn_api", "--sdk_location", "container" ])), server) beam_artifact_api_pb2_grpc.add_ArtifactRetrievalServiceServicer_to_server( artifact_service.ArtifactRetrievalService( artifact_service.BeamFilesystemHandler(None).file_reader), server) server.add_insecure_port('localhost:{}'.format(options.port)) server.start() _LOGGER.info('Listening for expansion requests at %d', options.port) signal.signal(signal.SIGTERM, cleanup) signal.signal(signal.SIGINT, cleanup) # blocking main thread forever. signal.pause()
def _stage_files(self, files): """Utility method to stage files. Args: files: a list of tuples of the form [(local_name, remote_name),...] describing the name of the artifacts in local temp folder and desired name in staging location. """ server = grpc.server(thread_pool_executor.shared_unbounded_instance()) staging_service = TestLocalFileSystemLegacyArtifactStagingServiceServicer( self._remote_dir) beam_artifact_api_pb2_grpc.add_LegacyArtifactStagingServiceServicer_to_server( staging_service, server) test_port = server.add_insecure_port('[::]:0') server.start() stager = portable_stager.PortableStager( artifact_service_channel=grpc.insecure_channel('localhost:%s' % test_port), staging_session_token='token') for from_file, to_file in files: stager.stage_artifact(local_path_to_artifact=os.path.join( self._temp_dir, from_file), artifact_name=to_file) stager.commit_manifest() return staging_service.manifest.artifact, staging_service.retrieval_tokens
def process_bundle( self, inputs, # type: Mapping[str, execution.PartitionableBuffer] expected_outputs, # type: DataOutput fired_timers, # type: Mapping[Tuple[str, str], execution.PartitionableBuffer] expected_output_timers, # type: Dict[Tuple[str, str], str] dry_run=False, ): # type: (...) -> BundleProcessResult part_inputs = [{} for _ in range(self._num_workers) ] # type: List[Dict[str, List[bytes]]] # Timers are only executed on the first worker # TODO(BEAM-9741): Split timers to multiple workers timer_inputs = [ fired_timers if i == 0 else {} for i in range(self._num_workers) ] for name, input in inputs.items(): for ix, part in enumerate(input.partition(self._num_workers)): part_inputs[ix][name] = part merged_result = None # type: Optional[beam_fn_api_pb2.InstructionResponse] split_result_list = [ ] # type: List[beam_fn_api_pb2.ProcessBundleSplitResponse] def execute(part_map_input_timers): # type: (...) -> BundleProcessResult part_map, input_timers = part_map_input_timers bundle_manager = BundleManager( self.bundle_context_manager, self._progress_frequency, cache_token_generator=self._cache_token_generator) return bundle_manager.process_bundle(part_map, expected_outputs, input_timers, expected_output_timers, dry_run) with thread_pool_executor.shared_unbounded_instance() as executor: for result, split_result in executor.map( execute, zip( part_inputs, # pylint: disable=zip-builtin-not-iterating timer_inputs)): split_result_list += split_result if merged_result is None: merged_result = result else: merged_result = beam_fn_api_pb2.InstructionResponse( process_bundle=beam_fn_api_pb2.ProcessBundleResponse( monitoring_infos=monitoring_infos.consolidate( itertools.chain( result.process_bundle.monitoring_infos, merged_result.process_bundle. monitoring_infos))), error=result.error or merged_result.error) assert merged_result is not None return merged_result, split_result_list
def setUp(self): self.num_request = 3 self.test_status_service = BeamFnStatusServicer(self.num_request) self.server = grpc.server( thread_pool_executor.shared_unbounded_instance()) beam_fn_api_pb2_grpc.add_BeamFnWorkerStatusServicer_to_server( self.test_status_service, self.server) self.test_port = self.server.add_insecure_port('[::]:0') self.server.start() self.url = 'localhost:%s' % self.test_port self.fn_status_handler = FnApiWorkerStatusHandler(self.url)
def start(self): if not self._worker_address: worker_server = grpc.server( thread_pool_executor.shared_unbounded_instance()) worker_address = 'localhost:%s' % worker_server.add_insecure_port( '[::]:0') beam_fn_api_pb2_grpc.add_BeamFnExternalWorkerPoolServicer_to_server( self, worker_server) worker_server.start() self._worker_address = worker_address atexit.register(functools.partial(worker_server.stop, 1)) return self._worker_address
def get_responses(self, instruction_requests): """Evaluates and returns {id: InstructionResponse} for the requests.""" test_controller = BeamFnControlServicer(instruction_requests) server = grpc.server(thread_pool_executor.shared_unbounded_instance()) beam_fn_api_pb2_grpc.add_BeamFnControlServicer_to_server( test_controller, server) test_port = server.add_insecure_port("[::]:0") server.start() harness = sdk_worker.SdkHarness("localhost:%s" % test_port, state_cache_size=100) harness.run() return test_controller.responses
def setUp(self): self.test_logging_service = BeamFnLoggingServicer() self.server = grpc.server(thread_pool_executor.shared_unbounded_instance()) beam_fn_api_pb2_grpc.add_BeamFnLoggingServicer_to_server( self.test_logging_service, self.server) self.test_port = self.server.add_insecure_port('[::]:0') self.server.start() self.logging_service_descriptor = endpoints_pb2.ApiServiceDescriptor() self.logging_service_descriptor.url = 'localhost:%s' % self.test_port self.fn_log_handler = log_handler.FnApiLogRecordHandler( self.logging_service_descriptor) logging.getLogger().setLevel(logging.INFO) logging.getLogger().addHandler(self.fn_log_handler)
def start_grpc_server(self, port=0): self._server = grpc.server( thread_pool_executor.shared_unbounded_instance()) port = self._server.add_insecure_port('%s:%d' % (self.get_bind_address(), port)) beam_job_api_pb2_grpc.add_JobServiceServicer_to_server( self, self._server) beam_artifact_api_pb2_grpc.add_ArtifactStagingServiceServicer_to_server( self._artifact_service, self._server) hostname = self.get_service_address() self._artifact_staging_endpoint = endpoints_pb2.ApiServiceDescriptor( url='%s:%d' % (hostname, port)) self._server.start() _LOGGER.info('Grpc server started at %s on port %d' % (hostname, port)) return port
def start_grpc_server(self, port=0): options = [("grpc.max_receive_message_length", -1), ("grpc.max_send_message_length", -1), ("grpc.http2.max_pings_without_data", 0), ("grpc.http2.max_ping_strikes", 0)] self._server = grpc.server( thread_pool_executor.shared_unbounded_instance(), options=options) port = self._server.add_insecure_port( '%s:%d' % (self.get_bind_address(), port)) beam_job_api_pb2_grpc.add_JobServiceServicer_to_server(self, self._server) beam_artifact_api_pb2_grpc.add_ArtifactStagingServiceServicer_to_server( self._artifact_service, self._server) hostname = self.get_service_address() self._artifact_staging_endpoint = endpoints_pb2.ApiServiceDescriptor( url='%s:%d' % (hostname, port)) self._server.start() _LOGGER.info('Grpc server started at %s on port %d' % (hostname, port)) return port
def run(self): options = [("grpc.http2.max_pings_without_data", 0), ("grpc.http2.max_ping_strikes", 0)] logging_server = grpc.server( thread_pool_executor.shared_unbounded_instance(), options=options) logging_port = logging_server.add_insecure_port('[::]:0') logging_server.start() logging_servicer = BeamFnLoggingServicer() beam_fn_api_pb2_grpc.add_BeamFnLoggingServicer_to_server( logging_servicer, logging_server) logging_descriptor = text_format.MessageToString( endpoints_pb2.ApiServiceDescriptor(url='localhost:%s' % logging_port)) control_descriptor = text_format.MessageToString( endpoints_pb2.ApiServiceDescriptor(url=self._control_address)) pipeline_options = json_format.MessageToJson( self._provision_info.provision_info.pipeline_options) env_dict = dict(os.environ, CONTROL_API_SERVICE_DESCRIPTOR=control_descriptor, LOGGING_API_SERVICE_DESCRIPTOR=logging_descriptor, PIPELINE_OPTIONS=pipeline_options) # only add worker_id when it is set. if self._worker_id: env_dict['WORKER_ID'] = self._worker_id with worker_handlers.SUBPROCESS_LOCK: p = subprocess.Popen(self._worker_command_line, shell=True, env=env_dict) try: p.wait() if p.returncode: raise RuntimeError( 'Worker subprocess exited with return code %s' % p.returncode) finally: if p.poll() is None: p.kill() logging_server.stop(0)
def main(unused_argv): parser = argparse.ArgumentParser() parser.add_argument('-p', '--port', type=int, help='port on which to serve the job api') options = parser.parse_args() global server server = grpc.server(thread_pool_executor.shared_unbounded_instance()) beam_expansion_api_pb2_grpc.add_ExpansionServiceServicer_to_server( expansion_service.ExpansionServiceServicer( PipelineOptions([ "--experiments", "beam_fn_api", "--sdk_location", "container" ])), server) server.add_insecure_port('localhost:{}'.format(options.port)) server.start() _LOGGER.info('Listening for expansion requests at %d', options.port) signal.signal(signal.SIGTERM, cleanup) signal.signal(signal.SIGINT, cleanup) # blocking main thread forever. signal.pause()
def _check_fn_registration_multi_request(self, *args): """Check the function registration calls to the sdk_harness. Args: tuple of request_count, number of process_bundles per request and workers counts to process the request. """ for (request_count, process_bundles_per_request) in args: requests = [] process_bundle_descriptors = [] for i in range(request_count): pbd = self._get_process_bundles(i, process_bundles_per_request) process_bundle_descriptors.extend(pbd) requests.append( beam_fn_api_pb2.InstructionRequest( instruction_id=str(i), register=beam_fn_api_pb2.RegisterRequest( process_bundle_descriptor=process_bundle_descriptors ))) test_controller = BeamFnControlServicer(requests) server = grpc.server( thread_pool_executor.shared_unbounded_instance()) beam_fn_api_pb2_grpc.add_BeamFnControlServicer_to_server( test_controller, server) test_port = server.add_insecure_port("[::]:0") server.start() harness = sdk_worker.SdkHarness("localhost:%s" % test_port, state_cache_size=100) harness.run() self.assertEqual( harness._bundle_processor_cache.fns, {item.id: item for item in process_bundle_descriptors})
def _grpc_data_channel_test(self, time_based_flush=False): if time_based_flush: data_servicer = data_plane.BeamFnDataServicer( data_buffer_time_limit_ms=100) else: data_servicer = data_plane.BeamFnDataServicer() worker_id = 'worker_0' data_channel_service = \ data_servicer.get_conn_by_worker_id(worker_id) server = grpc.server(thread_pool_executor.shared_unbounded_instance()) beam_fn_api_pb2_grpc.add_BeamFnDataServicer_to_server( data_servicer, server) test_port = server.add_insecure_port('[::]:0') server.start() grpc_channel = grpc.insecure_channel('localhost:%s' % test_port) # Add workerId to the grpc channel grpc_channel = grpc.intercept_channel(grpc_channel, WorkerIdInterceptor(worker_id)) data_channel_stub = beam_fn_api_pb2_grpc.BeamFnDataStub(grpc_channel) if time_based_flush: data_channel_client = data_plane.GrpcClientDataChannel( data_channel_stub, data_buffer_time_limit_ms=100) else: data_channel_client = data_plane.GrpcClientDataChannel( data_channel_stub) try: self._data_channel_test(data_channel_service, data_channel_client, time_based_flush) finally: data_channel_client.close() data_channel_service.close() data_channel_client.wait() data_channel_service.wait()
def __init__( self, state, # type: StateServicer provision_info, # type: Optional[ExtendedProvisionInfo] worker_manager, # type: WorkerHandlerManager ): # type: (...) -> None self.state = state self.provision_info = provision_info self.control_server = grpc.server( thread_pool_executor.shared_unbounded_instance()) self.control_port = self.control_server.add_insecure_port('[::]:0') self.control_address = 'localhost:%s' % self.control_port # Options to have no limits (-1) on the size of the messages # received or sent over the data plane. The actual buffer size # is controlled in a layer above. no_max_message_sizes = [("grpc.max_receive_message_length", -1), ("grpc.max_send_message_length", -1)] self.data_server = grpc.server( thread_pool_executor.shared_unbounded_instance(), options=no_max_message_sizes) self.data_port = self.data_server.add_insecure_port('[::]:0') self.state_server = grpc.server( thread_pool_executor.shared_unbounded_instance(), options=no_max_message_sizes) self.state_port = self.state_server.add_insecure_port('[::]:0') self.control_handler = BeamFnControlServicer(worker_manager) beam_fn_api_pb2_grpc.add_BeamFnControlServicer_to_server( self.control_handler, self.control_server) # If we have provision info, serve these off the control port as well. if self.provision_info: if self.provision_info.provision_info: beam_provision_api_pb2_grpc.add_ProvisionServiceServicer_to_server( BasicProvisionService(self.provision_info.provision_info, worker_manager), self.control_server) def open_uncompressed(f): # type: (str) -> BinaryIO return filesystems.FileSystems.open( f, compression_type=CompressionTypes.UNCOMPRESSED) beam_artifact_api_pb2_grpc.add_ArtifactRetrievalServiceServicer_to_server( artifact_service.ArtifactRetrievalService( file_reader=open_uncompressed), self.control_server) self.data_plane_handler = data_plane.BeamFnDataServicer( DATA_BUFFER_TIME_LIMIT_MS) beam_fn_api_pb2_grpc.add_BeamFnDataServicer_to_server( self.data_plane_handler, self.data_server) beam_fn_api_pb2_grpc.add_BeamFnStateServicer_to_server( GrpcStateServicer(state), self.state_server) self.logging_server = grpc.server( thread_pool_executor.shared_unbounded_instance(), options=no_max_message_sizes) self.logging_port = self.logging_server.add_insecure_port('[::]:0') beam_fn_api_pb2_grpc.add_BeamFnLoggingServicer_to_server( BasicLoggingService(), self.logging_server) _LOGGER.info('starting control server on port %s', self.control_port) _LOGGER.info('starting data server on port %s', self.data_port) _LOGGER.info('starting state server on port %s', self.state_port) _LOGGER.info('starting logging server on port %s', self.logging_port) self.logging_server.start() self.state_server.start() self.data_server.start() self.control_server.start()
def __init__( self, control_address, # type: str credentials=None, worker_id=None, # type: Optional[str] # Caching is disabled by default state_cache_size=0, # time-based data buffering is disabled by default data_buffer_time_limit_ms=0, profiler_factory=None, # type: Optional[Callable[..., Profile]] status_address=None, # type: Optional[str] ): self._alive = True self._worker_index = 0 self._worker_id = worker_id self._state_cache = StateCache(state_cache_size) if credentials is None: _LOGGER.info('Creating insecure control channel for %s.', control_address) self._control_channel = GRPCChannelFactory.insecure_channel( control_address) else: _LOGGER.info('Creating secure control channel for %s.', control_address) self._control_channel = GRPCChannelFactory.secure_channel( control_address, credentials) grpc.channel_ready_future(self._control_channel).result(timeout=60) _LOGGER.info('Control channel established.') self._control_channel = grpc.intercept_channel( self._control_channel, WorkerIdInterceptor(self._worker_id)) self._data_channel_factory = data_plane.GrpcClientDataChannelFactory( credentials, self._worker_id, data_buffer_time_limit_ms) self._state_handler_factory = GrpcStateHandlerFactory( self._state_cache, credentials) self._profiler_factory = profiler_factory self._fns = KeyedDefaultDict( lambda id: self._control_stub.GetProcessBundleDescriptor( beam_fn_api_pb2.GetProcessBundleDescriptorRequest( process_bundle_descriptor_id=id) )) # type: Mapping[str, beam_fn_api_pb2.ProcessBundleDescriptor] # BundleProcessor cache across all workers. self._bundle_processor_cache = BundleProcessorCache( state_handler_factory=self._state_handler_factory, data_channel_factory=self._data_channel_factory, fns=self._fns) if status_address: try: self._status_handler = FnApiWorkerStatusHandler( status_address, self._bundle_processor_cache ) # type: Optional[FnApiWorkerStatusHandler] except Exception: traceback_string = traceback.format_exc() _LOGGER.warning( 'Error creating worker status request handler, ' 'skipping status report. Trace back: %s' % traceback_string) else: self._status_handler = None # TODO(BEAM-8998) use common # thread_pool_executor.shared_unbounded_instance() to process bundle # progress once dataflow runner's excessive progress polling is removed. self._report_progress_executor = futures.ThreadPoolExecutor( max_workers=1) self._worker_thread_pool = thread_pool_executor.shared_unbounded_instance( ) self._responses = queue.Queue( ) # type: queue.Queue[beam_fn_api_pb2.InstructionResponse] _LOGGER.info( 'Initializing SDKHarness with unbounded number of workers.')