| # |
| # Licensed to the Apache Software Foundation (ASF) under one or more |
| # contributor license agreements. See the NOTICE file distributed with |
| # this work for additional information regarding copyright ownership. |
| # The ASF licenses this file to You under the Apache License, Version 2.0 |
| # (the "License"); you may not use this file except in compliance with |
| # the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # |
| |
| """A PipelineRunner using the SDK harness. |
| """ |
| from __future__ import absolute_import |
| from __future__ import print_function |
| |
| import collections |
| import contextlib |
| import copy |
| import itertools |
| import logging |
| import os |
| import queue |
| import subprocess |
| import sys |
| import threading |
| import time |
| import uuid |
| from builtins import object |
| from concurrent import futures |
| |
| import grpc |
| |
| import apache_beam as beam # pylint: disable=ungrouped-imports |
| from apache_beam import coders |
| from apache_beam import metrics |
| from apache_beam.coders.coder_impl import create_InputStream |
| from apache_beam.coders.coder_impl import create_OutputStream |
| from apache_beam.metrics import monitoring_infos |
| from apache_beam.metrics.execution import MetricKey |
| from apache_beam.metrics.execution import MetricsEnvironment |
| from apache_beam.metrics.metricbase import MetricName |
| from apache_beam.options import pipeline_options |
| from apache_beam.options.value_provider import RuntimeValueProvider |
| from apache_beam.portability import common_urns |
| from apache_beam.portability import python_urns |
| from apache_beam.portability.api import beam_artifact_api_pb2_grpc |
| from apache_beam.portability.api import beam_fn_api_pb2 |
| from apache_beam.portability.api import beam_fn_api_pb2_grpc |
| from apache_beam.portability.api import beam_provision_api_pb2 |
| from apache_beam.portability.api import beam_provision_api_pb2_grpc |
| from apache_beam.portability.api import beam_runner_api_pb2 |
| from apache_beam.portability.api import endpoints_pb2 |
| from apache_beam.runners import pipeline_context |
| from apache_beam.runners import runner |
| from apache_beam.runners.portability import artifact_service |
| from apache_beam.runners.portability import fn_api_runner_transforms |
| from apache_beam.runners.portability.fn_api_runner_transforms import create_buffer_id |
| from apache_beam.runners.portability.fn_api_runner_transforms import only_element |
| from apache_beam.runners.portability.fn_api_runner_transforms import split_buffer_id |
| from apache_beam.runners.portability.fn_api_runner_transforms import unique_name |
| from apache_beam.runners.worker import bundle_processor |
| from apache_beam.runners.worker import data_plane |
| from apache_beam.runners.worker import sdk_worker |
| from apache_beam.runners.worker.channel_factory import GRPCChannelFactory |
| from apache_beam.transforms import trigger |
| from apache_beam.transforms.window import GlobalWindows |
| from apache_beam.utils import profiler |
| from apache_beam.utils import proto_utils |
| |
| # This module is experimental. No backwards-compatibility guarantees. |
| |
| ENCODED_IMPULSE_VALUE = beam.coders.WindowedValueCoder( |
| beam.coders.BytesCoder(), |
| beam.coders.coders.GlobalWindowCoder()).get_impl().encode_nested( |
| beam.transforms.window.GlobalWindows.windowed_value(b'')) |
| |
| |
| class BeamFnControlServicer(beam_fn_api_pb2_grpc.BeamFnControlServicer): |
| |
| UNSTARTED_STATE = 'unstarted' |
| STARTED_STATE = 'started' |
| DONE_STATE = 'done' |
| |
| _DONE_MARKER = object() |
| |
| def __init__(self): |
| self._push_queue = queue.Queue() |
| self._futures_by_id = dict() |
| self._read_thread = threading.Thread( |
| name='beam_control_read', target=self._read) |
| self._uid_counter = 0 |
| self._state = self.UNSTARTED_STATE |
| self._lock = threading.Lock() |
| |
| def Control(self, iterator, context): |
| with self._lock: |
| if self._state == self.DONE_STATE: |
| return |
| else: |
| self._state = self.STARTED_STATE |
| self._inputs = iterator |
| # Note: We only support one client for now. |
| self._read_thread.start() |
| while True: |
| to_push = self._push_queue.get() |
| if to_push is self._DONE_MARKER: |
| return |
| yield to_push |
| |
| def _read(self): |
| for data in self._inputs: |
| self._futures_by_id.pop(data.instruction_id).set(data) |
| |
| def push(self, item): |
| if item is self._DONE_MARKER: |
| future = None |
| else: |
| if not item.instruction_id: |
| self._uid_counter += 1 |
| item.instruction_id = 'control_%s' % self._uid_counter |
| future = ControlFuture(item.instruction_id) |
| self._futures_by_id[item.instruction_id] = future |
| self._push_queue.put(item) |
| return future |
| |
| def done(self): |
| with self._lock: |
| if self._state == self.STARTED_STATE: |
| self.push(self._DONE_MARKER) |
| self._read_thread.join() |
| self._state = self.DONE_STATE |
| |
| |
| class _GroupingBuffer(object): |
| """Used to accumulate groupded (shuffled) results.""" |
| def __init__(self, pre_grouped_coder, post_grouped_coder, windowing): |
| self._key_coder = pre_grouped_coder.key_coder() |
| self._pre_grouped_coder = pre_grouped_coder |
| self._post_grouped_coder = post_grouped_coder |
| self._table = collections.defaultdict(list) |
| self._windowing = windowing |
| self._grouped_output = None |
| |
| def append(self, elements_data): |
| if self._grouped_output: |
| raise RuntimeError('Grouping table append after read.') |
| input_stream = create_InputStream(elements_data) |
| coder_impl = self._pre_grouped_coder.get_impl() |
| key_coder_impl = self._key_coder.get_impl() |
| # TODO(robertwb): We could optimize this even more by using a |
| # window-dropping coder for the data plane. |
| is_trivial_windowing = self._windowing.is_default() |
| while input_stream.size() > 0: |
| windowed_key_value = coder_impl.decode_from_stream(input_stream, True) |
| key, value = windowed_key_value.value |
| self._table[key_coder_impl.encode(key)].append( |
| value if is_trivial_windowing |
| else windowed_key_value.with_value(value)) |
| |
| def __iter__(self): |
| if not self._grouped_output: |
| output_stream = create_OutputStream() |
| if self._windowing.is_default(): |
| globally_window = GlobalWindows.windowed_value(None).with_value |
| windowed_key_values = lambda key, values: [ |
| globally_window((key, values))] |
| else: |
| trigger_driver = trigger.create_trigger_driver(self._windowing, True) |
| windowed_key_values = trigger_driver.process_entire_key |
| coder_impl = self._post_grouped_coder.get_impl() |
| key_coder_impl = self._key_coder.get_impl() |
| for encoded_key, windowed_values in self._table.items(): |
| key = key_coder_impl.decode(encoded_key) |
| for wkvs in windowed_key_values(key, windowed_values): |
| coder_impl.encode_to_stream(wkvs, output_stream, True) |
| self._grouped_output = [output_stream.get()] |
| self._table = None |
| return iter(self._grouped_output) |
| |
| |
| class _WindowGroupingBuffer(object): |
| """Used to partition windowed side inputs.""" |
| def __init__(self, access_pattern, coder): |
| # Here's where we would use a different type of partitioning |
| # (e.g. also by key) for a different access pattern. |
| if access_pattern.urn == common_urns.side_inputs.ITERABLE.urn: |
| self._kv_extrator = lambda value: ('', value) |
| self._key_coder = coders.SingletonCoder('') |
| self._value_coder = coder.wrapped_value_coder |
| elif access_pattern.urn == common_urns.side_inputs.MULTIMAP.urn: |
| self._kv_extrator = lambda value: value |
| self._key_coder = coder.wrapped_value_coder.key_coder() |
| self._value_coder = ( |
| coder.wrapped_value_coder.value_coder()) |
| else: |
| raise ValueError( |
| "Unknown access pattern: '%s'" % access_pattern.urn) |
| self._windowed_value_coder = coder |
| self._window_coder = coder.window_coder |
| self._values_by_window = collections.defaultdict(list) |
| |
| def append(self, elements_data): |
| input_stream = create_InputStream(elements_data) |
| while input_stream.size() > 0: |
| windowed_value = self._windowed_value_coder.get_impl( |
| ).decode_from_stream(input_stream, True) |
| key, value = self._kv_extrator(windowed_value.value) |
| for window in windowed_value.windows: |
| self._values_by_window[key, window].append(value) |
| |
| def encoded_items(self): |
| value_coder_impl = self._value_coder.get_impl() |
| key_coder_impl = self._key_coder.get_impl() |
| for (key, window), values in self._values_by_window.items(): |
| encoded_window = self._window_coder.encode(window) |
| encoded_key = key_coder_impl.encode_nested(key) |
| output_stream = create_OutputStream() |
| for value in values: |
| value_coder_impl.encode_to_stream(value, output_stream, True) |
| yield encoded_key, encoded_window, output_stream.get() |
| |
| |
| class FnApiRunner(runner.PipelineRunner): |
| |
| def __init__( |
| self, |
| default_environment=None, |
| bundle_repeat=0, |
| use_state_iterables=False, |
| provision_info=None): |
| """Creates a new Fn API Runner. |
| |
| Args: |
| default_environment: the default environment to use for UserFns. |
| bundle_repeat: replay every bundle this many extra times, for profiling |
| and debugging |
| use_state_iterables: Intentionally split gbk iterables over state API |
| (for testing) |
| provision_info: provisioning info to make available to workers, or None |
| """ |
| super(FnApiRunner, self).__init__() |
| self._last_uid = -1 |
| self._default_environment = ( |
| default_environment |
| or beam_runner_api_pb2.Environment(urn=python_urns.EMBEDDED_PYTHON)) |
| self._bundle_repeat = bundle_repeat |
| self._progress_frequency = None |
| self._profiler_factory = None |
| self._use_state_iterables = use_state_iterables |
| self._provision_info = provision_info |
| |
| def _next_uid(self): |
| self._last_uid += 1 |
| return str(self._last_uid) |
| |
| def run_pipeline(self, pipeline, options): |
| MetricsEnvironment.set_metrics_supported(False) |
| RuntimeValueProvider.set_runtime_options({}) |
| # This is sometimes needed if type checking is disabled |
| # to enforce that the inputs (and outputs) of GroupByKey operations |
| # are known to be KVs. |
| from apache_beam.runners.dataflow.dataflow_runner import DataflowRunner |
| # TODO: Move group_by_key_input_visitor() to a non-dataflow specific file. |
| pipeline.visit(DataflowRunner.group_by_key_input_visitor()) |
| self._bundle_repeat = self._bundle_repeat or options.view_as( |
| pipeline_options.DirectOptions).direct_runner_bundle_repeat |
| self._profiler_factory = profiler.Profile.factory_from_options( |
| options.view_as(pipeline_options.ProfilingOptions)) |
| self._latest_run_result = self.run_via_runner_api(pipeline.to_runner_api( |
| default_environment=self._default_environment)) |
| return self._latest_run_result |
| |
| def run_via_runner_api(self, pipeline_proto): |
| return self.run_stages(*self.create_stages(pipeline_proto)) |
| |
| @contextlib.contextmanager |
| def maybe_profile(self): |
| if self._profiler_factory: |
| try: |
| profile_id = 'direct-' + subprocess.check_output( |
| ['git', 'rev-parse', '--abbrev-ref', 'HEAD'] |
| ).decode(errors='ignore').strip() |
| except subprocess.CalledProcessError: |
| profile_id = 'direct-unknown' |
| profiler = self._profiler_factory(profile_id, time_prefix='') |
| else: |
| profiler = None |
| |
| if profiler: |
| with profiler: |
| yield |
| if not self._bundle_repeat: |
| logging.warning( |
| 'The --direct_runner_bundle_repeat option is not set; ' |
| 'a significant portion of the profile may be one-time overhead.') |
| path = profiler.profile_output |
| print('CPU Profile written to %s' % path) |
| try: |
| import gprof2dot # pylint: disable=unused-variable |
| if not subprocess.call([ |
| sys.executable, '-m', 'gprof2dot', |
| '-f', 'pstats', path, '-o', path + '.dot']): |
| if not subprocess.call( |
| ['dot', '-Tsvg', '-o', path + '.svg', path + '.dot']): |
| print('CPU Profile rendering at file://%s.svg' |
| % os.path.abspath(path)) |
| except ImportError: |
| # pylint: disable=superfluous-parens |
| print('Please install gprof2dot and dot for profile renderings.') |
| |
| else: |
| # Empty context. |
| yield |
| |
| def create_stages(self, pipeline_proto): |
| return fn_api_runner_transforms.create_and_optimize_stages( |
| copy.deepcopy(pipeline_proto), |
| phases=[fn_api_runner_transforms.annotate_downstream_side_inputs, |
| fn_api_runner_transforms.fix_side_input_pcoll_coders, |
| fn_api_runner_transforms.lift_combiners, |
| fn_api_runner_transforms.expand_sdf, |
| fn_api_runner_transforms.expand_gbk, |
| fn_api_runner_transforms.sink_flattens, |
| fn_api_runner_transforms.greedily_fuse, |
| fn_api_runner_transforms.read_to_impulse, |
| fn_api_runner_transforms.impulse_to_input, |
| fn_api_runner_transforms.inject_timer_pcollections, |
| fn_api_runner_transforms.sort_stages, |
| fn_api_runner_transforms.window_pcollection_coders], |
| known_runner_urns=frozenset([ |
| common_urns.primitives.FLATTEN.urn, |
| common_urns.primitives.GROUP_BY_KEY.urn]), |
| use_state_iterables=self._use_state_iterables) |
| |
| def run_stages(self, stage_context, stages): |
| worker_handler_manager = WorkerHandlerManager( |
| stage_context.components.environments, self._provision_info) |
| metrics_by_stage = {} |
| monitoring_infos_by_stage = {} |
| |
| try: |
| with self.maybe_profile(): |
| pcoll_buffers = collections.defaultdict(list) |
| for stage in stages: |
| stage_results = self.run_stage( |
| worker_handler_manager.get_worker_handler, |
| stage_context.components, |
| stage, |
| pcoll_buffers, |
| stage_context.safe_coders) |
| metrics_by_stage[stage.name] = stage_results.process_bundle.metrics |
| monitoring_infos_by_stage[stage.name] = ( |
| stage_results.process_bundle.monitoring_infos) |
| finally: |
| worker_handler_manager.close_all() |
| return RunnerResult( |
| runner.PipelineState.DONE, monitoring_infos_by_stage, metrics_by_stage) |
| |
| def run_stage(self, |
| worker_handler_factory, |
| pipeline_components, |
| stage, |
| pcoll_buffers, |
| safe_coders): |
| """Run an individual stage. |
| |
| Args: |
| worker_handler_factory: A ``callable`` that takes in an environment, and |
| returns a ``WorkerHandler`` class. |
| pipeline_components: TODO |
| stage: TODO |
| pcoll_buffers: TODO |
| safe_coders: TODO |
| """ |
| |
| def iterable_state_write(values, element_coder_impl): |
| token = unique_name(None, 'iter').encode('ascii') |
| out = create_OutputStream() |
| for element in values: |
| element_coder_impl.encode_to_stream(element, out, True) |
| controller.state.blocking_append( |
| beam_fn_api_pb2.StateKey( |
| runner=beam_fn_api_pb2.StateKey.Runner(key=token)), |
| out.get()) |
| return token |
| |
| controller = worker_handler_factory(stage.environment) |
| context = pipeline_context.PipelineContext( |
| pipeline_components, iterable_state_write=iterable_state_write) |
| data_api_service_descriptor = controller.data_api_service_descriptor() |
| |
| def extract_endpoints(stage): |
| # Returns maps of transform names to PCollection identifiers. |
| # Also mutates IO stages to point to the data ApiServiceDescriptor. |
| data_input = {} |
| data_side_input = {} |
| data_output = {} |
| for transform in stage.transforms: |
| if transform.spec.urn in (bundle_processor.DATA_INPUT_URN, |
| bundle_processor.DATA_OUTPUT_URN): |
| pcoll_id = transform.spec.payload |
| if transform.spec.urn == bundle_processor.DATA_INPUT_URN: |
| target = transform.unique_name, only_element(transform.outputs) |
| if pcoll_id == fn_api_runner_transforms.IMPULSE_BUFFER: |
| data_input[target] = [ENCODED_IMPULSE_VALUE] |
| else: |
| data_input[target] = pcoll_buffers[pcoll_id] |
| coder_id = pipeline_components.pcollections[ |
| only_element(transform.outputs.values())].coder_id |
| elif transform.spec.urn == bundle_processor.DATA_OUTPUT_URN: |
| target = transform.unique_name, only_element(transform.inputs) |
| data_output[target] = pcoll_id |
| coder_id = pipeline_components.pcollections[ |
| only_element(transform.inputs.values())].coder_id |
| else: |
| raise NotImplementedError |
| data_spec = beam_fn_api_pb2.RemoteGrpcPort(coder_id=coder_id) |
| if data_api_service_descriptor: |
| data_spec.api_service_descriptor.url = ( |
| data_api_service_descriptor.url) |
| transform.spec.payload = data_spec.SerializeToString() |
| elif transform.spec.urn in fn_api_runner_transforms.PAR_DO_URNS: |
| payload = proto_utils.parse_Bytes( |
| transform.spec.payload, beam_runner_api_pb2.ParDoPayload) |
| for tag, si in payload.side_inputs.items(): |
| data_side_input[transform.unique_name, tag] = ( |
| create_buffer_id(transform.inputs[tag]), si.access_pattern) |
| return data_input, data_side_input, data_output |
| |
| logging.info('Running %s', stage.name) |
| logging.debug(' %s', stage) |
| data_input, data_side_input, data_output = extract_endpoints(stage) |
| |
| process_bundle_descriptor = beam_fn_api_pb2.ProcessBundleDescriptor( |
| id=self._next_uid(), |
| transforms={transform.unique_name: transform |
| for transform in stage.transforms}, |
| pcollections=dict(pipeline_components.pcollections.items()), |
| coders=dict(pipeline_components.coders.items()), |
| windowing_strategies=dict( |
| pipeline_components.windowing_strategies.items()), |
| environments=dict(pipeline_components.environments.items())) |
| |
| if controller.state_api_service_descriptor(): |
| process_bundle_descriptor.state_api_service_descriptor.url = ( |
| controller.state_api_service_descriptor().url) |
| |
| # Store the required side inputs into state. |
| for (transform_id, tag), (buffer_id, si) in data_side_input.items(): |
| _, pcoll_id = split_buffer_id(buffer_id) |
| value_coder = context.coders[safe_coders[ |
| pipeline_components.pcollections[pcoll_id].coder_id]] |
| elements_by_window = _WindowGroupingBuffer(si, value_coder) |
| for element_data in pcoll_buffers[buffer_id]: |
| elements_by_window.append(element_data) |
| for key, window, elements_data in elements_by_window.encoded_items(): |
| state_key = beam_fn_api_pb2.StateKey( |
| multimap_side_input=beam_fn_api_pb2.StateKey.MultimapSideInput( |
| ptransform_id=transform_id, |
| side_input_id=tag, |
| window=window, |
| key=key)) |
| controller.state.blocking_append(state_key, elements_data) |
| |
| def get_buffer(buffer_id): |
| kind, name = split_buffer_id(buffer_id) |
| if kind in ('materialize', 'timers'): |
| if buffer_id not in pcoll_buffers: |
| # Just store the data chunks for replay. |
| pcoll_buffers[buffer_id] = list() |
| elif kind == 'group': |
| # This is a grouping write, create a grouping buffer if needed. |
| if buffer_id not in pcoll_buffers: |
| original_gbk_transform = name |
| transform_proto = pipeline_components.transforms[ |
| original_gbk_transform] |
| input_pcoll = only_element(list(transform_proto.inputs.values())) |
| output_pcoll = only_element(list(transform_proto.outputs.values())) |
| pre_gbk_coder = context.coders[safe_coders[ |
| pipeline_components.pcollections[input_pcoll].coder_id]] |
| post_gbk_coder = context.coders[safe_coders[ |
| pipeline_components.pcollections[output_pcoll].coder_id]] |
| windowing_strategy = context.windowing_strategies[ |
| pipeline_components |
| .pcollections[output_pcoll].windowing_strategy_id] |
| pcoll_buffers[buffer_id] = _GroupingBuffer( |
| pre_gbk_coder, post_gbk_coder, windowing_strategy) |
| else: |
| # These should be the only two identifiers we produce for now, |
| # but special side input writes may go here. |
| raise NotImplementedError(buffer_id) |
| return pcoll_buffers[buffer_id] |
| |
| def get_input_coder_impl(transform_id): |
| return context.coders[safe_coders[ |
| beam_fn_api_pb2.RemoteGrpcPort.FromString( |
| process_bundle_descriptor.transforms[transform_id].spec.payload |
| ).coder_id |
| ]].get_impl() |
| |
| for k in range(self._bundle_repeat): |
| try: |
| controller.state.checkpoint() |
| BundleManager( |
| controller, lambda pcoll_id: [], get_input_coder_impl, |
| process_bundle_descriptor, self._progress_frequency, k |
| ).process_bundle(data_input, data_output) |
| finally: |
| controller.state.restore() |
| |
| result, splits = BundleManager( |
| controller, get_buffer, get_input_coder_impl, process_bundle_descriptor, |
| self._progress_frequency).process_bundle( |
| data_input, data_output) |
| |
| def input_for(ptransform_id, input_id): |
| input_pcoll = process_bundle_descriptor.transforms[ |
| ptransform_id].inputs[input_id] |
| for read_id, proto in process_bundle_descriptor.transforms.items(): |
| if (proto.spec.urn == bundle_processor.DATA_INPUT_URN |
| and input_pcoll in proto.outputs.values()): |
| return read_id, 'out' |
| raise RuntimeError( |
| 'No IO transform feeds %s' % ptransform_id) |
| |
| last_result = result |
| last_sent = data_input |
| |
| while True: |
| deferred_inputs = collections.defaultdict(list) |
| for transform_id, timer_writes in stage.timer_pcollections: |
| |
| # Queue any set timers as new inputs. |
| windowed_timer_coder_impl = context.coders[ |
| pipeline_components.pcollections[timer_writes].coder_id].get_impl() |
| written_timers = get_buffer( |
| create_buffer_id(timer_writes, kind='timers')) |
| if written_timers: |
| # Keep only the "last" timer set per key and window. |
| timers_by_key_and_window = {} |
| for elements_data in written_timers: |
| input_stream = create_InputStream(elements_data) |
| while input_stream.size() > 0: |
| windowed_key_timer = windowed_timer_coder_impl.decode_from_stream( |
| input_stream, True) |
| key, _ = windowed_key_timer.value |
| # TODO: Explode and merge windows. |
| assert len(windowed_key_timer.windows) == 1 |
| timers_by_key_and_window[ |
| key, windowed_key_timer.windows[0]] = windowed_key_timer |
| out = create_OutputStream() |
| for windowed_key_timer in timers_by_key_and_window.values(): |
| windowed_timer_coder_impl.encode_to_stream( |
| windowed_key_timer, out, True) |
| deferred_inputs[transform_id, 'out'] = [out.get()] |
| written_timers[:] = [] |
| |
| # Queue any process-initiated delayed bundle applications. |
| for delayed_application in last_result.process_bundle.residual_roots: |
| deferred_inputs[ |
| input_for( |
| delayed_application.application.ptransform_id, |
| delayed_application.application.input_id) |
| ].append(delayed_application.application.element) |
| |
| # Queue any runner-initiated delayed bundle applications. |
| prev_stops = {} |
| for split in splits: |
| for delayed_application in split.residual_roots: |
| deferred_inputs[ |
| input_for( |
| delayed_application.application.ptransform_id, |
| delayed_application.application.input_id) |
| ].append(delayed_application.application.element) |
| for channel_split in split.channel_splits: |
| coder_impl = get_input_coder_impl(channel_split.ptransform_id) |
| # TODO(SDF): This requires determanistic ordering of buffer iteration. |
| # TODO(SDF): The return split is in terms of indices. Ideally, |
| # a runner could map these back to actual positions to effectively |
| # describe the two "halves" of the now-split range. Even if we have |
| # to buffer each element we send (or at the very least a bit of |
| # metadata, like position, about each of them) this should be doable |
| # if they're already in memory and we are bounding the buffer size |
| # (e.g. to 10mb plus whatever is eagerly read from the SDK). In the |
| # case of non-split-points, we can either immediately replay the |
| # "non-split-position" elements or record them as we do the other |
| # delayed applications. |
| |
| # Decode and recode to split the encoded buffer by element index. |
| all_elements = list(coder_impl.decode_all(b''.join(last_sent[ |
| channel_split.ptransform_id, channel_split.input_id]))) |
| residual_elements = all_elements[ |
| channel_split.first_residual_element : prev_stops.get( |
| channel_split.ptransform_id, len(all_elements)) + 1] |
| if residual_elements: |
| deferred_inputs[ |
| channel_split.ptransform_id, channel_split.input_id].append( |
| coder_impl.encode_all(residual_elements)) |
| prev_stops[ |
| channel_split.ptransform_id] = channel_split.last_primary_element |
| |
| if deferred_inputs: |
| # The worker will be waiting on these inputs as well. |
| for other_input in data_input: |
| if other_input not in deferred_inputs: |
| deferred_inputs[other_input] = [] |
| # TODO(robertwb): merge results |
| last_result, splits = BundleManager( |
| controller, |
| get_buffer, |
| get_input_coder_impl, |
| process_bundle_descriptor, |
| self._progress_frequency, |
| True).process_bundle(deferred_inputs, data_output) |
| last_sent = deferred_inputs |
| 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, |
| last_result.process_bundle.monitoring_infos))), |
| error=result.error or last_result.error) |
| else: |
| break |
| |
| return result |
| |
| # These classes are used to interact with the worker. |
| |
| class StateServicer(beam_fn_api_pb2_grpc.BeamFnStateServicer): |
| |
| class CopyOnWriteState(object): |
| def __init__(self, underlying): |
| self._underlying = underlying |
| self._overlay = {} |
| |
| def __getitem__(self, key): |
| if key in self._overlay: |
| return self._overlay[key] |
| else: |
| return FnApiRunner.StateServicer.CopyOnWriteList( |
| self._underlying, self._overlay, key) |
| |
| def __delitem__(self, key): |
| self._overlay[key] = [] |
| |
| def commit(self): |
| self._underlying.update(self._overlay) |
| return self._underlying |
| |
| class CopyOnWriteList(object): |
| def __init__(self, underlying, overlay, key): |
| self._underlying = underlying |
| self._overlay = overlay |
| self._key = key |
| |
| def __iter__(self): |
| if self._key in self._overlay: |
| return iter(self._overlay[self._key]) |
| else: |
| return iter(self._underlying[self._key]) |
| |
| def append(self, item): |
| if self._key not in self._overlay: |
| self._overlay[self._key] = list(self._underlying[self._key]) |
| self._overlay[self._key].append(item) |
| |
| def __init__(self): |
| self._lock = threading.Lock() |
| self._state = collections.defaultdict(list) |
| self._checkpoint = None |
| self._use_continuation_tokens = False |
| self._continuations = {} |
| |
| def checkpoint(self): |
| assert self._checkpoint is None |
| self._checkpoint = self._state |
| self._state = FnApiRunner.StateServicer.CopyOnWriteState(self._state) |
| |
| def commit(self): |
| self._state.commit() |
| self._state = self._checkpoint.commit() |
| self._checkpoint = None |
| |
| def restore(self): |
| self._state = self._checkpoint |
| self._checkpoint = None |
| |
| @contextlib.contextmanager |
| def process_instruction_id(self, unused_instruction_id): |
| yield |
| |
| def blocking_get(self, state_key, continuation_token=None): |
| with self._lock: |
| full_state = self._state[self._to_key(state_key)] |
| if self._use_continuation_tokens: |
| # The token is "nonce:index". |
| if not continuation_token: |
| token_base = 'token_%x' % len(self._continuations) |
| self._continuations[token_base] = tuple(full_state) |
| return b'', '%s:0' % token_base |
| else: |
| token_base, index = continuation_token.split(':') |
| ix = int(index) |
| full_state = self._continuations[token_base] |
| if ix == len(full_state): |
| return b'', None |
| else: |
| return full_state[ix], '%s:%d' % (token_base, ix + 1) |
| else: |
| assert not continuation_token |
| return b''.join(full_state), None |
| |
| def blocking_append(self, state_key, data): |
| with self._lock: |
| self._state[self._to_key(state_key)].append(data) |
| |
| def blocking_clear(self, state_key): |
| with self._lock: |
| del self._state[self._to_key(state_key)] |
| |
| @staticmethod |
| def _to_key(state_key): |
| return state_key.SerializeToString() |
| |
| class GrpcStateServicer(beam_fn_api_pb2_grpc.BeamFnStateServicer): |
| def __init__(self, state): |
| self._state = state |
| |
| def State(self, request_stream, context=None): |
| # Note that this eagerly mutates state, assuming any failures are fatal. |
| # Thus it is safe to ignore instruction_reference. |
| for request in request_stream: |
| request_type = request.WhichOneof('request') |
| if request_type == 'get': |
| data, continuation_token = self._state.blocking_get( |
| request.state_key, request.get.continuation_token) |
| yield beam_fn_api_pb2.StateResponse( |
| id=request.id, |
| get=beam_fn_api_pb2.StateGetResponse( |
| data=data, continuation_token=continuation_token)) |
| elif request_type == 'append': |
| self._state.blocking_append(request.state_key, request.append.data) |
| yield beam_fn_api_pb2.StateResponse( |
| id=request.id, |
| append=beam_fn_api_pb2.StateAppendResponse()) |
| elif request_type == 'clear': |
| self._state.blocking_clear(request.state_key) |
| yield beam_fn_api_pb2.StateResponse( |
| id=request.id, |
| clear=beam_fn_api_pb2.StateClearResponse()) |
| else: |
| raise NotImplementedError('Unknown state request: %s' % request_type) |
| |
| class SingletonStateHandlerFactory(sdk_worker.StateHandlerFactory): |
| """A singleton cache for a StateServicer.""" |
| |
| def __init__(self, state_handler): |
| self._state_handler = state_handler |
| |
| def create_state_handler(self, api_service_descriptor): |
| """Returns the singleton state handler.""" |
| return self._state_handler |
| |
| def close(self): |
| """Does nothing.""" |
| pass |
| |
| |
| class WorkerHandler(object): |
| |
| _registered_environments = {} |
| |
| def __init__( |
| self, control_handler, data_plane_handler, state, provision_info): |
| self.control_handler = control_handler |
| self.data_plane_handler = data_plane_handler |
| self.state = state |
| self.provision_info = provision_info |
| |
| def close(self): |
| self.stop_worker() |
| |
| def start_worker(self): |
| raise NotImplementedError |
| |
| def stop_worker(self): |
| raise NotImplementedError |
| |
| def data_api_service_descriptor(self): |
| raise NotImplementedError |
| |
| def state_api_service_descriptor(self): |
| raise NotImplementedError |
| |
| def logging_api_service_descriptor(self): |
| raise NotImplementedError |
| |
| @classmethod |
| def register_environment(cls, urn, payload_type): |
| def wrapper(constructor): |
| cls._registered_environments[urn] = constructor, payload_type |
| return constructor |
| return wrapper |
| |
| @classmethod |
| def create(cls, environment, state, provision_info): |
| constructor, payload_type = cls._registered_environments[environment.urn] |
| return constructor( |
| proto_utils.parse_Bytes(environment.payload, payload_type), |
| state, |
| provision_info) |
| |
| |
| @WorkerHandler.register_environment(python_urns.EMBEDDED_PYTHON, None) |
| class EmbeddedWorkerHandler(WorkerHandler): |
| """An in-memory controller for fn API control, state and data planes.""" |
| |
| def __init__(self, unused_payload, state, provision_info): |
| super(EmbeddedWorkerHandler, self).__init__( |
| self, data_plane.InMemoryDataChannel(), state, provision_info) |
| self.worker = sdk_worker.SdkWorker( |
| sdk_worker.BundleProcessorCache( |
| FnApiRunner.SingletonStateHandlerFactory(self.state), |
| data_plane.InMemoryDataChannelFactory( |
| self.data_plane_handler.inverse()), |
| {})) |
| self._uid_counter = 0 |
| |
| def push(self, request): |
| if not request.instruction_id: |
| self._uid_counter += 1 |
| request.instruction_id = 'control_%s' % self._uid_counter |
| logging.debug('CONTROL REQUEST %s', request) |
| response = self.worker.do_instruction(request) |
| logging.debug('CONTROL RESPONSE %s', response) |
| return ControlFuture(request.instruction_id, response) |
| |
| def start_worker(self): |
| pass |
| |
| def stop_worker(self): |
| self.worker.stop() |
| |
| def done(self): |
| pass |
| |
| def data_api_service_descriptor(self): |
| return None |
| |
| def state_api_service_descriptor(self): |
| return None |
| |
| def logging_api_service_descriptor(self): |
| return None |
| |
| |
| class BasicLoggingService(beam_fn_api_pb2_grpc.BeamFnLoggingServicer): |
| |
| LOG_LEVEL_MAP = { |
| beam_fn_api_pb2.LogEntry.Severity.CRITICAL: logging.CRITICAL, |
| beam_fn_api_pb2.LogEntry.Severity.ERROR: logging.ERROR, |
| beam_fn_api_pb2.LogEntry.Severity.WARN: logging.WARNING, |
| beam_fn_api_pb2.LogEntry.Severity.NOTICE: logging.INFO + 1, |
| beam_fn_api_pb2.LogEntry.Severity.INFO: logging.INFO, |
| beam_fn_api_pb2.LogEntry.Severity.DEBUG: logging.DEBUG, |
| beam_fn_api_pb2.LogEntry.Severity.TRACE: logging.DEBUG - 1, |
| beam_fn_api_pb2.LogEntry.Severity.UNSPECIFIED: logging.NOTSET, |
| } |
| |
| def Logging(self, log_messages, context=None): |
| yield beam_fn_api_pb2.LogControl() |
| for log_message in log_messages: |
| for log in log_message.log_entries: |
| logging.log(self.LOG_LEVEL_MAP[log.severity], str(log)) |
| |
| |
| class BasicProvisionService( |
| beam_provision_api_pb2_grpc.ProvisionServiceServicer): |
| |
| def __init__(self, info): |
| self._info = info |
| |
| def GetProvisionInfo(self, request, context=None): |
| return beam_provision_api_pb2.GetProvisionInfoResponse( |
| info=self._info) |
| |
| |
| class GrpcWorkerHandler(WorkerHandler): |
| """An grpc based controller for fn API control, state and data planes.""" |
| |
| _DEFAULT_SHUTDOWN_TIMEOUT_SECS = 5 |
| |
| def __init__(self, state, provision_info): |
| self.state = state |
| self.provision_info = provision_info |
| self.control_server = grpc.server( |
| futures.ThreadPoolExecutor(max_workers=10)) |
| 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( |
| futures.ThreadPoolExecutor(max_workers=10), |
| options=no_max_message_sizes) |
| self.data_port = self.data_server.add_insecure_port('[::]:0') |
| |
| self.state_server = grpc.server( |
| futures.ThreadPoolExecutor(max_workers=10), |
| options=no_max_message_sizes) |
| self.state_port = self.state_server.add_insecure_port('[::]:0') |
| |
| self.control_handler = BeamFnControlServicer() |
| 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: |
| provision_info = self.provision_info.provision_info |
| if not provision_info.worker_id: |
| provision_info = copy.copy(provision_info) |
| provision_info.worker_id = str(uuid.uuid4()) |
| beam_provision_api_pb2_grpc.add_ProvisionServiceServicer_to_server( |
| BasicProvisionService(self.provision_info.provision_info), |
| self.control_server) |
| |
| if self.provision_info.artifact_staging_dir: |
| m = beam_artifact_api_pb2_grpc |
| m.add_ArtifactRetrievalServiceServicer_to_server( |
| artifact_service.BeamFilesystemArtifactService( |
| self.provision_info.artifact_staging_dir), |
| self.control_server) |
| |
| self.data_plane_handler = data_plane.GrpcServerDataChannel() |
| 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( |
| FnApiRunner.GrpcStateServicer(state), |
| self.state_server) |
| |
| self.logging_server = grpc.server( |
| futures.ThreadPoolExecutor(max_workers=2), |
| 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) |
| |
| logging.info('starting control server on port %s', self.control_port) |
| logging.info('starting data server on port %s', self.data_port) |
| logging.info('starting state server on port %s', self.state_port) |
| logging.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 data_api_service_descriptor(self): |
| return endpoints_pb2.ApiServiceDescriptor( |
| url='localhost:%s' % self.data_port) |
| |
| def state_api_service_descriptor(self): |
| return endpoints_pb2.ApiServiceDescriptor( |
| url='localhost:%s' % self.state_port) |
| |
| def logging_api_service_descriptor(self): |
| return endpoints_pb2.ApiServiceDescriptor( |
| url='localhost:%s' % self.logging_port) |
| |
| def close(self): |
| self.control_handler.done() |
| self.data_plane_handler.close() |
| to_wait = [ |
| self.control_server.stop(self._DEFAULT_SHUTDOWN_TIMEOUT_SECS), |
| self.data_server.stop(self._DEFAULT_SHUTDOWN_TIMEOUT_SECS), |
| self.state_server.stop(self._DEFAULT_SHUTDOWN_TIMEOUT_SECS), |
| self.logging_server.stop(self._DEFAULT_SHUTDOWN_TIMEOUT_SECS) |
| ] |
| for w in to_wait: |
| w.wait() |
| super(GrpcWorkerHandler, self).close() |
| |
| |
| @WorkerHandler.register_environment( |
| common_urns.environments.EXTERNAL.urn, beam_runner_api_pb2.ExternalPayload) |
| class ExternalWorkerHandler(GrpcWorkerHandler): |
| def __init__(self, external_payload, state, provision_info): |
| super(ExternalWorkerHandler, self).__init__(state, provision_info) |
| self._external_payload = external_payload |
| |
| def start_worker(self): |
| stub = beam_fn_api_pb2_grpc.BeamFnExternalWorkerPoolStub( |
| GRPCChannelFactory.insecure_channel( |
| self._external_payload.endpoint.url)) |
| response = stub.NotifyRunnerAvailable( |
| beam_fn_api_pb2.NotifyRunnerAvailableRequest( |
| worker_id='worker_%s' % uuid.uuid4(), |
| control_endpoint=endpoints_pb2.ApiServiceDescriptor( |
| url=self.control_address), |
| logging_endpoint=self.logging_api_service_descriptor(), |
| params=self._external_payload.params)) |
| if response.error: |
| raise RuntimeError("Error starting worker: %s" % response.error) |
| |
| def stop_worker(self): |
| pass |
| |
| |
| @WorkerHandler.register_environment(python_urns.EMBEDDED_PYTHON_GRPC, bytes) |
| class EmbeddedGrpcWorkerHandler(GrpcWorkerHandler): |
| def __init__(self, num_workers_payload, state, provision_info): |
| super(EmbeddedGrpcWorkerHandler, self).__init__(state, provision_info) |
| self._num_threads = int(num_workers_payload) if num_workers_payload else 1 |
| |
| def start_worker(self): |
| self.worker = sdk_worker.SdkHarness( |
| self.control_address, worker_count=self._num_threads) |
| self.worker_thread = threading.Thread( |
| name='run_worker', target=self.worker.run) |
| self.worker_thread.daemon = True |
| self.worker_thread.start() |
| |
| def stop_worker(self): |
| self.worker_thread.join() |
| |
| |
| @WorkerHandler.register_environment(python_urns.SUBPROCESS_SDK, bytes) |
| class SubprocessSdkWorkerHandler(GrpcWorkerHandler): |
| def __init__(self, worker_command_line, state, provision_info): |
| super(SubprocessSdkWorkerHandler, self).__init__(state, provision_info) |
| self._worker_command_line = worker_command_line |
| |
| def start_worker(self): |
| from apache_beam.runners.portability import local_job_service |
| self.worker = local_job_service.SubprocessSdkWorker( |
| self._worker_command_line, self.control_address) |
| self.worker_thread = threading.Thread( |
| name='run_worker', target=self.worker.run) |
| self.worker_thread.start() |
| |
| def stop_worker(self): |
| self.worker_thread.join() |
| |
| |
| @WorkerHandler.register_environment(common_urns.environments.DOCKER.urn, |
| beam_runner_api_pb2.DockerPayload) |
| class DockerSdkWorkerHandler(GrpcWorkerHandler): |
| def __init__(self, payload, state, provision_info): |
| super(DockerSdkWorkerHandler, self).__init__(state, provision_info) |
| self._container_image = payload.container_image |
| self._container_id = None |
| |
| def start_worker(self): |
| try: |
| subprocess.check_call(['docker', 'pull', self._container_image]) |
| except Exception: |
| logging.info('Unable to pull image %s' % self._container_image) |
| self._container_id = subprocess.check_output( |
| ['docker', |
| 'run', |
| '-d', |
| # TODO: credentials |
| '--network=host', |
| self._container_image, |
| '--id=%s' % uuid.uuid4(), |
| '--logging_endpoint=%s' % self.logging_api_service_descriptor().url, |
| '--control_endpoint=%s' % self.control_address, |
| '--artifact_endpoint=%s' % self.control_address, |
| '--provision_endpoint=%s' % self.control_address, |
| ]).strip() |
| while True: |
| logging.info('Waiting for docker to start up...') |
| status = subprocess.check_output([ |
| 'docker', |
| 'inspect', |
| '-f', |
| '{{.State.Status}}', |
| self._container_id]).strip() |
| if status == 'running': |
| break |
| elif status in ('dead', 'exited'): |
| subprocess.call([ |
| 'docker', |
| 'container', |
| 'logs', |
| self._container_id]) |
| raise RuntimeError('SDK failed to start.') |
| time.sleep(1) |
| |
| def stop_worker(self): |
| if self._container_id: |
| subprocess.call([ |
| 'docker', |
| 'kill', |
| self._container_id]) |
| |
| |
| class WorkerHandlerManager(object): |
| def __init__(self, environments, job_provision_info=None): |
| self._environments = environments |
| self._job_provision_info = job_provision_info |
| self._cached_handlers = {} |
| self._state = FnApiRunner.StateServicer() # rename? |
| |
| def get_worker_handler(self, environment_id): |
| if environment_id is None: |
| # Any environment will do, pick one arbitrarily. |
| environment_id = next(iter(self._environments.keys())) |
| environment = self._environments[environment_id] |
| |
| worker_handler = self._cached_handlers.get(environment_id) |
| if worker_handler is None: |
| worker_handler = self._cached_handlers[ |
| environment_id] = WorkerHandler.create( |
| environment, self._state, self._job_provision_info) |
| worker_handler.start_worker() |
| return worker_handler |
| |
| def close_all(self): |
| for controller in set(self._cached_handlers.values()): |
| try: |
| controller.close() |
| except Exception: |
| logging.error("Error closing controller %s" % controller, exc_info=True) |
| self._cached_handlers = {} |
| |
| |
| class ExtendedProvisionInfo(object): |
| def __init__(self, provision_info=None, artifact_staging_dir=None): |
| self.provision_info = ( |
| provision_info or beam_provision_api_pb2.ProvisionInfo()) |
| self.artifact_staging_dir = artifact_staging_dir |
| |
| |
| _split_managers = [] |
| |
| |
| @contextlib.contextmanager |
| def split_manager(stage_name, split_manager): |
| """Registers a split manager to control the flow of elements to a given stage. |
| |
| Used for testing. |
| |
| A split manager should be a coroutine yielding desired split fractions, |
| receiving the corresponding split results. Currently, only one input is |
| supported. |
| """ |
| try: |
| _split_managers.append((stage_name, split_manager)) |
| yield |
| finally: |
| _split_managers.pop() |
| |
| |
| class BundleManager(object): |
| |
| _uid_counter = 0 |
| |
| def __init__( |
| self, controller, get_buffer, get_input_coder_impl, bundle_descriptor, |
| progress_frequency=None, skip_registration=False): |
| self._controller = controller |
| self._get_buffer = get_buffer |
| self._get_input_coder_impl = get_input_coder_impl |
| self._bundle_descriptor = bundle_descriptor |
| self._registered = skip_registration |
| self._progress_frequency = progress_frequency |
| |
| def process_bundle(self, inputs, expected_outputs): |
| # Unique id for the instruction processing this bundle. |
| BundleManager._uid_counter += 1 |
| process_bundle_id = 'bundle_%s' % BundleManager._uid_counter |
| |
| # Register the bundle descriptor, if needed. |
| if self._registered: |
| registration_future = None |
| else: |
| process_bundle_registration = beam_fn_api_pb2.InstructionRequest( |
| register=beam_fn_api_pb2.RegisterRequest( |
| process_bundle_descriptor=[self._bundle_descriptor])) |
| registration_future = self._controller.control_handler.push( |
| process_bundle_registration) |
| self._registered = True |
| |
| unique_names = set( |
| t.unique_name for t in self._bundle_descriptor.transforms.values()) |
| for stage_name, candidate in reversed(_split_managers): |
| if (stage_name in unique_names |
| or (stage_name + '/Process') in unique_names): |
| split_manager = candidate |
| break |
| else: |
| split_manager = None |
| |
| if not split_manager: |
| # Write all the input data to the channel immediately. |
| for (transform_id, name), elements in inputs.items(): |
| data_out = self._controller.data_plane_handler.output_stream( |
| process_bundle_id, beam_fn_api_pb2.Target( |
| primitive_transform_reference=transform_id, name=name)) |
| for element_data in elements: |
| data_out.write(element_data) |
| data_out.close() |
| |
| split_results = [] |
| |
| # Actually start the bundle. |
| if registration_future and registration_future.get().error: |
| raise RuntimeError(registration_future.get().error) |
| process_bundle = beam_fn_api_pb2.InstructionRequest( |
| instruction_id=process_bundle_id, |
| process_bundle=beam_fn_api_pb2.ProcessBundleRequest( |
| process_bundle_descriptor_reference=self._bundle_descriptor.id)) |
| result_future = self._controller.control_handler.push(process_bundle) |
| |
| with ProgressRequester( |
| self._controller, process_bundle_id, self._progress_frequency): |
| if split_manager: |
| (read_transform_id, name), buffer_data = only_element(inputs.items()) |
| num_elements = len(list( |
| self._get_input_coder_impl(read_transform_id).decode_all( |
| b''.join(buffer_data)))) |
| |
| # Start the split manager in case it wants to set any breakpoints. |
| split_manager_generator = split_manager(num_elements) |
| try: |
| split_fraction = next(split_manager_generator) |
| done = False |
| except StopIteration: |
| done = True |
| |
| # Send all the data. |
| data_out = self._controller.data_plane_handler.output_stream( |
| process_bundle_id, |
| beam_fn_api_pb2.Target( |
| primitive_transform_reference=read_transform_id, name=name)) |
| data_out.write(b''.join(buffer_data)) |
| data_out.close() |
| |
| # Execute the requested splits. |
| while not done: |
| if split_fraction is None: |
| split_result = None |
| else: |
| split_request = beam_fn_api_pb2.InstructionRequest( |
| process_bundle_split= |
| beam_fn_api_pb2.ProcessBundleSplitRequest( |
| instruction_reference=process_bundle_id, |
| desired_splits={ |
| read_transform_id: |
| beam_fn_api_pb2.ProcessBundleSplitRequest.DesiredSplit( |
| fraction_of_remainder=split_fraction, |
| estimated_input_elements=num_elements) |
| })) |
| split_response = self._controller.control_handler.push( |
| split_request).get() |
| for t in (0.05, 0.1, 0.2): |
| waiting = ('Instruction not running', 'not yet scheduled') |
| if any(msg in split_response.error for msg in waiting): |
| time.sleep(t) |
| split_response = self._controller.control_handler.push( |
| split_request).get() |
| if 'Unknown process bundle' in split_response.error: |
| # It may have finished too fast. |
| split_result = None |
| elif split_response.error: |
| raise RuntimeError(split_response.error) |
| else: |
| split_result = split_response.process_bundle_split |
| split_results.append(split_result) |
| try: |
| split_fraction = split_manager_generator.send(split_result) |
| except StopIteration: |
| break |
| |
| # Gather all output data. |
| expected_targets = [ |
| beam_fn_api_pb2.Target(primitive_transform_reference=transform_id, |
| name=output_name) |
| for (transform_id, output_name), _ in expected_outputs.items()] |
| logging.debug('Gather all output data from %s.', expected_targets) |
| for output in self._controller.data_plane_handler.input_elements( |
| process_bundle_id, |
| expected_targets, |
| abort_callback=lambda: (result_future.is_done() |
| and result_future.get().error)): |
| target_tuple = ( |
| output.target.primitive_transform_reference, output.target.name) |
| if target_tuple in expected_outputs: |
| self._get_buffer(expected_outputs[target_tuple]).append(output.data) |
| |
| logging.debug('Wait for the bundle to finish.') |
| result = result_future.get() |
| |
| if result.error: |
| raise RuntimeError(result.error) |
| |
| if result.process_bundle.requires_finalization: |
| finalize_request = beam_fn_api_pb2.InstructionRequest( |
| finalize_bundle= |
| beam_fn_api_pb2.FinalizeBundleRequest( |
| instruction_reference=process_bundle_id |
| )) |
| self._controller.control_handler.push( |
| finalize_request) |
| |
| return result, split_results |
| |
| |
| class ProgressRequester(threading.Thread): |
| def __init__(self, controller, instruction_id, frequency, callback=None): |
| super(ProgressRequester, self).__init__() |
| self._controller = controller |
| self._instruction_id = instruction_id |
| self._frequency = frequency |
| self._done = False |
| self._latest_progress = None |
| self._callback = callback |
| self.daemon = True |
| |
| def __enter__(self): |
| if self._frequency: |
| self.start() |
| |
| def __exit__(self, *unused_exc_info): |
| if self._frequency: |
| self.stop() |
| |
| def run(self): |
| while not self._done: |
| try: |
| progress_result = self._controller.control_handler.push( |
| beam_fn_api_pb2.InstructionRequest( |
| process_bundle_progress= |
| beam_fn_api_pb2.ProcessBundleProgressRequest( |
| instruction_reference=self._instruction_id))).get() |
| self._latest_progress = progress_result.process_bundle_progress |
| if self._callback: |
| self._callback(self._latest_progress) |
| except Exception as exn: |
| logging.error("Bad progress: %s", exn) |
| time.sleep(self._frequency) |
| |
| def stop(self): |
| self._done = True |
| |
| |
| class ControlFuture(object): |
| def __init__(self, instruction_id, response=None): |
| self.instruction_id = instruction_id |
| if response: |
| self._response = response |
| else: |
| self._response = None |
| self._condition = threading.Condition() |
| |
| def is_done(self): |
| return self._response is not None |
| |
| def set(self, response): |
| with self._condition: |
| self._response = response |
| self._condition.notify_all() |
| |
| def get(self, timeout=None): |
| if not self._response: |
| with self._condition: |
| if not self._response: |
| self._condition.wait(timeout) |
| return self._response |
| |
| |
| class FnApiMetrics(metrics.metric.MetricResults): |
| def __init__(self, step_monitoring_infos, user_metrics_only=True): |
| """Used for querying metrics from the PipelineResult object. |
| |
| step_monitoring_infos: Per step metrics specified as MonitoringInfos. |
| use_monitoring_infos: If true, return the metrics based on the |
| step_monitoring_infos. |
| """ |
| self._counters = {} |
| self._distributions = {} |
| self._gauges = {} |
| self._user_metrics_only = user_metrics_only |
| self._init_metrics_from_monitoring_infos(step_monitoring_infos) |
| self._monitoring_infos = step_monitoring_infos |
| |
| def _init_metrics_from_monitoring_infos(self, step_monitoring_infos): |
| for smi in step_monitoring_infos.values(): |
| # Only include user metrics. |
| for mi in smi: |
| if (self._user_metrics_only and |
| not monitoring_infos.is_user_monitoring_info(mi)): |
| continue |
| key = self._to_metric_key(mi) |
| if monitoring_infos.is_counter(mi): |
| self._counters[key] = ( |
| monitoring_infos.extract_metric_result_map_value(mi)) |
| elif monitoring_infos.is_distribution(mi): |
| self._distributions[key] = ( |
| monitoring_infos.extract_metric_result_map_value(mi)) |
| elif monitoring_infos.is_gauge(mi): |
| self._gauges[key] = ( |
| monitoring_infos.extract_metric_result_map_value(mi)) |
| |
| def _to_metric_key(self, monitoring_info): |
| # Right now this assumes that all metrics have a PTRANSFORM |
| ptransform_id = monitoring_info.labels['PTRANSFORM'] |
| namespace, name = monitoring_infos.parse_namespace_and_name(monitoring_info) |
| return MetricKey(ptransform_id, MetricName(namespace, name)) |
| |
| def query(self, filter=None): |
| counters = [metrics.execution.MetricResult(k, v, v) |
| for k, v in self._counters.items() |
| if self.matches(filter, k)] |
| distributions = [metrics.execution.MetricResult(k, v, v) |
| for k, v in self._distributions.items() |
| if self.matches(filter, k)] |
| gauges = [metrics.execution.MetricResult(k, v, v) |
| for k, v in self._gauges.items() |
| if self.matches(filter, k)] |
| |
| return {self.COUNTERS: counters, |
| self.DISTRIBUTIONS: distributions, |
| self.GAUGES: gauges} |
| |
| def monitoring_infos(self): |
| return [item for sublist in self._monitoring_infos.values() for item in |
| sublist] |
| |
| |
| class RunnerResult(runner.PipelineResult): |
| def __init__(self, state, monitoring_infos_by_stage, metrics_by_stage): |
| super(RunnerResult, self).__init__(state) |
| self._monitoring_infos_by_stage = monitoring_infos_by_stage |
| self._metrics_by_stage = metrics_by_stage |
| self._metrics = None |
| self._monitoring_metrics = None |
| |
| def wait_until_finish(self, duration=None): |
| return self._state |
| |
| def metrics(self): |
| """Returns a queryable object including user metrics only.""" |
| if self._metrics is None: |
| self._metrics = FnApiMetrics( |
| self._monitoring_infos_by_stage, user_metrics_only=True) |
| return self._metrics |
| |
| def monitoring_metrics(self): |
| """Returns a queryable object including all metrics.""" |
| if self._monitoring_metrics is None: |
| self._monitoring_metrics = FnApiMetrics( |
| self._monitoring_infos_by_stage, user_metrics_only=False) |
| return self._monitoring_metrics |
