Source code for apache_beam.runners.direct.executor

#
# 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.
#

"""An executor that schedules and executes applied ptransforms."""

# pytype: skip-file

import collections
import itertools
import logging
import queue
import threading
import traceback
from typing import TYPE_CHECKING
from typing import Any
from typing import Dict
from typing import FrozenSet
from typing import Optional
from typing import Set
from weakref import WeakValueDictionary

from apache_beam.metrics.execution import MetricsContainer
from apache_beam.runners.worker import statesampler
from apache_beam.transforms import sideinputs
from apache_beam.utils import counters

if TYPE_CHECKING:
  from apache_beam import pvalue
  from apache_beam.runners.direct.bundle_factory import _Bundle
  from apache_beam.runners.direct.evaluation_context import EvaluationContext
  from apache_beam.runners.direct.transform_evaluator import TransformEvaluatorRegistry

_LOGGER = logging.getLogger(__name__)


class _ExecutorService(object):
  """Thread pool for executing tasks in parallel."""
  class CallableTask(object):
    def call(self, state_sampler):
      pass

    @property
    def name(self):
      return None

  class _ExecutorServiceWorker(threading.Thread):
    """Worker thread for executing a single task at a time."""

    # Amount to block waiting for getting an item from the queue in seconds.
    TIMEOUT = 5

    def __init__(
        self, queue: 'queue.Queue[_ExecutorService.CallableTask]', index):
      super().__init__()
      self.queue = queue
      self._index = index
      self._default_name = 'ExecutorServiceWorker-' + str(index)
      self._update_name()
      self.shutdown_requested = False

      # Stop worker thread when main thread exits.
      self.daemon = True
      self.start()

    def _update_name(self, task=None):
      if task and task.name:
        name = task.name
      else:
        name = self._default_name
      self.name = 'Thread: %d, %s (%s)' % (
          self._index, name, 'executing' if task else 'idle')

    def _get_task_or_none(self) -> Optional['_ExecutorService.CallableTask']:
      try:
        # Do not block indefinitely, otherwise we may not act for a requested
        # shutdown.
        return self.queue.get(
            timeout=_ExecutorService._ExecutorServiceWorker.TIMEOUT)
      except queue.Empty:
        return None

    def run(self):
      state_sampler = statesampler.StateSampler('', counters.CounterFactory())
      statesampler.set_current_tracker(state_sampler)
      while not self.shutdown_requested:
        task = self._get_task_or_none()
        if task:
          try:
            if not self.shutdown_requested:
              self._update_name(task)
              task.call(state_sampler)
              self._update_name()
          finally:
            self.queue.task_done()

    def shutdown(self):
      self.shutdown_requested = True

  def __init__(self, num_workers):
    self.queue: queue.Queue[_ExecutorService.CallableTask] = queue.Queue()
    self.workers = [
        _ExecutorService._ExecutorServiceWorker(self.queue, i)
        for i in range(num_workers)
    ]
    self.shutdown_requested = False

  def submit(self, task: '_ExecutorService.CallableTask') -> None:
    assert isinstance(task, _ExecutorService.CallableTask)
    if not self.shutdown_requested:
      self.queue.put(task)

  def await_completion(self):
    for worker in self.workers:
      worker.join()

  def shutdown(self):
    self.shutdown_requested = True

    for worker in self.workers:
      worker.shutdown()

    # Consume all the remaining items in the queue
    while not self.queue.empty():
      try:
        self.queue.get_nowait()
        self.queue.task_done()
      except queue.Empty:
        continue
    # All existing threads will eventually terminate (after they complete their
    # last task).


class _TransformEvaluationState(object):
  def __init__(self, executor_service, scheduled: Set['TransformExecutor']):
    self.executor_service = executor_service
    self.scheduled = scheduled

  def schedule(self, work):
    self.scheduled.add(work)
    self.executor_service.submit(work)

  def complete(self, completed_work):
    self.scheduled.remove(completed_work)


class _ParallelEvaluationState(_TransformEvaluationState):
  """A TransformEvaluationState with unlimited parallelism.

  Any TransformExecutor scheduled will be immediately submitted to the
  ExecutorService.

  A principal use of this is for evaluators that can generate output bundles
  only using the input bundle (e.g. ParDo).
  """
  pass


class _SerialEvaluationState(_TransformEvaluationState):
  """A TransformEvaluationState with a single work queue.

  Any TransformExecutor scheduled will be placed on the work queue. Only one
  item of work will be submitted to the ExecutorService at any time.

  A principal use of this is for evaluators that keeps a global state such as
  _GroupByKeyOnly.
  """
  def __init__(self, executor_service, scheduled):
    super().__init__(executor_service, scheduled)
    self.serial_queue = collections.deque()
    self.currently_evaluating = None
    self._lock = threading.Lock()

  def complete(self, completed_work):
    self._update_currently_evaluating(None, completed_work)
    super().complete(completed_work)

  def schedule(self, new_work):
    self._update_currently_evaluating(new_work, None)

  def _update_currently_evaluating(self, new_work, completed_work):
    with self._lock:
      if new_work:
        self.serial_queue.append(new_work)
      if completed_work:
        assert self.currently_evaluating == completed_work
        self.currently_evaluating = None
      if self.serial_queue and not self.currently_evaluating:
        next_work = self.serial_queue.pop()
        self.currently_evaluating = next_work
        super().schedule(next_work)


class _TransformExecutorServices(object):
  """Schedules and completes TransformExecutors.

  Controls the concurrency as appropriate for the applied transform the executor
  exists for.
  """
  def __init__(self, executor_service: _ExecutorService) -> None:
    self._executor_service = executor_service
    self._scheduled: Set[TransformExecutor] = set()
    self._parallel = _ParallelEvaluationState(
        self._executor_service, self._scheduled)
    self._serial_cache: WeakValueDictionary[
        Any, _SerialEvaluationState] = WeakValueDictionary()

  def parallel(self) -> _ParallelEvaluationState:
    return self._parallel

  def serial(self, step: Any) -> _SerialEvaluationState:
    cached = self._serial_cache.get(step)
    if not cached:
      cached = _SerialEvaluationState(self._executor_service, self._scheduled)
      self._serial_cache[step] = cached
    return cached

  @property
  def executors(self) -> FrozenSet['TransformExecutor']:
    return frozenset(self._scheduled)


class _CompletionCallback(object):
  """The default completion callback.

  The default completion callback is used to complete transform evaluations
  that are triggered due to the arrival of elements from an upstream transform,
  or for a source transform.
  """
  def __init__(
      self,
      evaluation_context: 'EvaluationContext',
      all_updates,
      timer_firings=None):
    self._evaluation_context = evaluation_context
    self._all_updates = all_updates
    self._timer_firings = timer_firings or []

  def handle_result(
      self, transform_executor, input_committed_bundle, transform_result):
    output_committed_bundles = self._evaluation_context.handle_result(
        input_committed_bundle, self._timer_firings, transform_result)
    for output_committed_bundle in output_committed_bundles:
      self._all_updates.offer(
          _ExecutorServiceParallelExecutor._ExecutorUpdate(
              transform_executor, committed_bundle=output_committed_bundle))
    for unprocessed_bundle in transform_result.unprocessed_bundles:
      self._all_updates.offer(
          _ExecutorServiceParallelExecutor._ExecutorUpdate(
              transform_executor, unprocessed_bundle=unprocessed_bundle))
    return output_committed_bundles

  def handle_exception(self, transform_executor, exception):
    self._all_updates.offer(
        _ExecutorServiceParallelExecutor._ExecutorUpdate(
            transform_executor, exception=exception))


[docs] class TransformExecutor(_ExecutorService.CallableTask): """For internal use only; no backwards-compatibility guarantees. TransformExecutor will evaluate a bundle using an applied ptransform. A CallableTask responsible for constructing a TransformEvaluator and evaluating it on some bundle of input, and registering the result using the completion callback. """ _MAX_RETRY_PER_BUNDLE = 4 def __init__( self, transform_evaluator_registry: 'TransformEvaluatorRegistry', evaluation_context: 'EvaluationContext', input_bundle: '_Bundle', fired_timers, applied_ptransform, completion_callback, transform_evaluation_state: _TransformEvaluationState): self._transform_evaluator_registry = transform_evaluator_registry self._evaluation_context = evaluation_context self._input_bundle = input_bundle # For non-empty bundles, store the window of the max EOW. # TODO(mariagh): Move to class _Bundle's inner _StackedWindowedValues self._latest_main_input_window = None if input_bundle.has_elements(): self._latest_main_input_window = input_bundle._elements[0].windows[0] for elem in input_bundle.get_elements_iterable(): if elem.windows[0].end > self._latest_main_input_window.end: self._latest_main_input_window = elem.windows[0] self._fired_timers = fired_timers self._applied_ptransform = applied_ptransform self._completion_callback = completion_callback self._transform_evaluation_state = transform_evaluation_state self._side_input_values: Dict[pvalue.AsSideInput, Any] = {} self.blocked = False self._call_count = 0 self._retry_count = 0 self._max_retries_per_bundle = TransformExecutor._MAX_RETRY_PER_BUNDLE
[docs] def call(self, state_sampler): self._call_count += 1 assert self._call_count <= (1 + len(self._applied_ptransform.side_inputs)) metrics_container = MetricsContainer(self._applied_ptransform.full_label) start_state = state_sampler.scoped_state( self._applied_ptransform.full_label, 'start', metrics_container=metrics_container) process_state = state_sampler.scoped_state( self._applied_ptransform.full_label, 'process', metrics_container=metrics_container) finish_state = state_sampler.scoped_state( self._applied_ptransform.full_label, 'finish', metrics_container=metrics_container) with start_state: # Side input initialization should be accounted for in start_state. for side_input in self._applied_ptransform.side_inputs: # Find the projection of main's window onto the side input's window. window_mapping_fn = side_input._view_options().get( 'window_mapping_fn', sideinputs._global_window_mapping_fn) main_onto_side_window = window_mapping_fn( self._latest_main_input_window) block_until = main_onto_side_window.end if side_input not in self._side_input_values: value = self._evaluation_context.get_value_or_block_until_ready( side_input, self, block_until) if not value: # Monitor task will reschedule this executor once the side input is # available. return self._side_input_values[side_input] = value side_input_values = [ self._side_input_values[side_input] for side_input in self._applied_ptransform.side_inputs ] while self._retry_count < self._max_retries_per_bundle: try: self.attempt_call( metrics_container, side_input_values, start_state, process_state, finish_state) break except Exception as e: self._retry_count += 1 _LOGGER.error( 'Exception at bundle %r, due to an exception.\n %s', self._input_bundle, traceback.format_exc()) if self._retry_count == self._max_retries_per_bundle: _LOGGER.error( 'Giving up after %s attempts.', self._max_retries_per_bundle) self._completion_callback.handle_exception(self, e) self._evaluation_context.metrics().commit_physical( self._input_bundle, metrics_container.get_cumulative()) self._transform_evaluation_state.complete(self)
[docs] def attempt_call( self, metrics_container, side_input_values, start_state, process_state, finish_state): """Attempts to run a bundle.""" evaluator = self._transform_evaluator_registry.get_evaluator( self._applied_ptransform, self._input_bundle, side_input_values) with start_state: evaluator.start_bundle() with process_state: if self._fired_timers: for timer_firing in self._fired_timers: evaluator.process_timer_wrapper(timer_firing) if self._input_bundle: for value in self._input_bundle.get_elements_iterable(): evaluator.process_element(value) with finish_state: result = evaluator.finish_bundle() result.logical_metric_updates = metrics_container.get_cumulative() self._completion_callback.handle_result(self, self._input_bundle, result) return result
[docs] class Executor(object): """For internal use only; no backwards-compatibility guarantees.""" def __init__(self, *args, **kwargs): self._executor = _ExecutorServiceParallelExecutor(*args, **kwargs)
[docs] def start(self, roots): self._executor.start(roots)
[docs] def await_completion(self): self._executor.await_completion()
[docs] def shutdown(self): self._executor.request_shutdown()
class _ExecutorServiceParallelExecutor(object): """An internal implementation for Executor.""" NUM_WORKERS = 1 def __init__( self, value_to_consumers, transform_evaluator_registry, evaluation_context: 'EvaluationContext'): self.executor_service = _ExecutorService( _ExecutorServiceParallelExecutor.NUM_WORKERS) self.transform_executor_services = _TransformExecutorServices( self.executor_service) self.value_to_consumers = value_to_consumers self.transform_evaluator_registry = transform_evaluator_registry self.evaluation_context = evaluation_context self.all_updates = _ExecutorServiceParallelExecutor._TypedUpdateQueue( _ExecutorServiceParallelExecutor._ExecutorUpdate) self.visible_updates = _ExecutorServiceParallelExecutor._TypedUpdateQueue( _ExecutorServiceParallelExecutor._VisibleExecutorUpdate) self.default_completion_callback = _CompletionCallback( evaluation_context, self.all_updates) def start(self, roots): self.root_nodes = frozenset(roots) self.all_nodes = frozenset( itertools.chain( roots, *itertools.chain(self.value_to_consumers.values()))) self.node_to_pending_bundles = {} for root_node in self.root_nodes: provider = ( self.transform_evaluator_registry.get_root_bundle_provider(root_node)) self.node_to_pending_bundles[root_node] = provider.get_root_bundles() self.executor_service.submit( _ExecutorServiceParallelExecutor._MonitorTask(self)) def await_completion(self): update = self.visible_updates.take() try: if update.exception: raise update.exception finally: self.executor_service.shutdown() self.executor_service.await_completion() def request_shutdown(self): self.executor_service.shutdown() self.executor_service.await_completion() self.evaluation_context.shutdown() def schedule_consumers(self, committed_bundle: '_Bundle') -> None: if committed_bundle.pcollection in self.value_to_consumers: consumers = self.value_to_consumers[committed_bundle.pcollection] for applied_ptransform in consumers: self.schedule_consumption( applied_ptransform, committed_bundle, [], self.default_completion_callback) def schedule_unprocessed_bundle(self, applied_ptransform, unprocessed_bundle): self.node_to_pending_bundles[applied_ptransform].append(unprocessed_bundle) def schedule_consumption( self, consumer_applied_ptransform, committed_bundle: '_Bundle', fired_timers, on_complete): """Schedules evaluation of the given bundle with the transform.""" assert consumer_applied_ptransform assert committed_bundle assert on_complete if self.transform_evaluator_registry.should_execute_serially( consumer_applied_ptransform): transform_executor_service: _TransformEvaluationState = ( self.transform_executor_services.serial(consumer_applied_ptransform)) else: transform_executor_service = self.transform_executor_services.parallel() transform_executor = TransformExecutor( self.transform_evaluator_registry, self.evaluation_context, committed_bundle, fired_timers, consumer_applied_ptransform, on_complete, transform_executor_service) transform_executor_service.schedule(transform_executor) class _TypedUpdateQueue(object): """Type checking update queue with blocking and non-blocking operations.""" def __init__(self, item_type): self._item_type = item_type self._queue = queue.Queue() def poll(self): try: item = self._queue.get_nowait() self._queue.task_done() return item except queue.Empty: return None def take(self): # The implementation of Queue.Queue.get() does not propagate # KeyboardInterrupts when a timeout is not used. We therefore use a # one-second timeout in the following loop to allow KeyboardInterrupts # to be correctly propagated. while True: try: item = self._queue.get(timeout=1) self._queue.task_done() return item except queue.Empty: pass def offer(self, item): assert isinstance(item, self._item_type) self._queue.put_nowait(item) class _ExecutorUpdate(object): """An internal status update on the state of the executor.""" def __init__( self, transform_executor, committed_bundle=None, unprocessed_bundle=None, exception=None): self.transform_executor = transform_executor # Exactly one of them should be not-None assert sum( [bool(committed_bundle), bool(unprocessed_bundle), bool(exception)]) == 1 self.committed_bundle = committed_bundle self.unprocessed_bundle = unprocessed_bundle self.exception = exception class _VisibleExecutorUpdate(object): """An update of interest to the user. Used for awaiting the completion to decide whether to return normally or raise an exception. """ def __init__(self, exception=None): self.finished = exception is not None self.exception = exception class _MonitorTask(_ExecutorService.CallableTask): """MonitorTask continuously runs to ensure that pipeline makes progress.""" def __init__(self, executor: '_ExecutorServiceParallelExecutor') -> None: self._executor = executor @property def name(self): return 'monitor' def call(self, state_sampler): try: update = self._executor.all_updates.poll() while update: if update.committed_bundle: self._executor.schedule_consumers(update.committed_bundle) elif update.unprocessed_bundle: self._executor.schedule_unprocessed_bundle( update.transform_executor._applied_ptransform, update.unprocessed_bundle) else: assert update.exception _LOGGER.warning( 'A task failed with exception: %s', update.exception) self._executor.visible_updates.offer( _ExecutorServiceParallelExecutor._VisibleExecutorUpdate( update.exception)) update = self._executor.all_updates.poll() self._executor.evaluation_context.schedule_pending_unblocked_tasks( self._executor.executor_service) self._add_work_if_necessary(self._fire_timers()) except Exception as e: # pylint: disable=broad-except _LOGGER.error('Monitor task died due to exception.\n %s', e) self._executor.visible_updates.offer( _ExecutorServiceParallelExecutor._VisibleExecutorUpdate(e)) finally: if not self._should_shutdown(): self._executor.executor_service.submit(self) def _should_shutdown(self) -> bool: """Checks whether the pipeline is completed and should be shut down. If there is anything in the queue of tasks to do or if there are any realtime timers set, do not shut down. Otherwise, check if all the transforms' watermarks are complete. If they are not, the pipeline is not progressing (stall detected). Whether the pipeline has stalled or not, the executor should shut down the pipeline. Returns: True only if the pipeline has reached a terminal state and should be shut down. """ if self._is_executing(): # There are some bundles still in progress. return False watermark_manager = self._executor.evaluation_context._watermark_manager _, any_unfired_realtime_timers = watermark_manager.extract_all_timers() if any_unfired_realtime_timers: return False else: if self._executor.evaluation_context.is_done(): self._executor.visible_updates.offer( _ExecutorServiceParallelExecutor._VisibleExecutorUpdate()) else: # Nothing is scheduled for execution, but watermarks incomplete. self._executor.visible_updates.offer( _ExecutorServiceParallelExecutor._VisibleExecutorUpdate(( Exception('Monitor task detected a pipeline stall.'), None, None))) self._executor.executor_service.shutdown() return True def _fire_timers(self): """Schedules triggered consumers if any timers fired. Returns: True if timers fired. """ transform_fired_timers, _ = ( self._executor.evaluation_context.extract_all_timers()) for applied_ptransform, fired_timers in transform_fired_timers: # Use an empty committed bundle. just to trigger. empty_bundle = ( self._executor.evaluation_context.create_empty_committed_bundle( applied_ptransform.inputs[0])) timer_completion_callback = _CompletionCallback( self._executor.evaluation_context, self._executor.all_updates, timer_firings=fired_timers) self._executor.schedule_consumption( applied_ptransform, empty_bundle, fired_timers, timer_completion_callback) return bool(transform_fired_timers) def _is_executing(self) -> bool: """Checks whether the job is still executing. Returns: True if there is at least one non-blocked TransformExecutor active.""" executors = self._executor.transform_executor_services.executors if not executors: # Nothing is executing. return False # Ensure that at least one of those executors is not blocked. for transform_executor in executors: if not transform_executor.blocked: return True return False def _add_work_if_necessary(self, timers_fired): """Adds more work from the roots if pipeline requires more input. If all active TransformExecutors are in a blocked state, add more work from root nodes that may have additional work. This ensures that if a pipeline has elements available from the root nodes it will add those elements when necessary. Args: timers_fired: True if any timers fired prior to this call. """ # If any timers have fired, they will add more work; No need to add more. if timers_fired: return if self._is_executing(): # We have at least one executor that can proceed without adding # additional work. return # All current TransformExecutors are blocked; add more work from any # pending bundles. for applied_ptransform in self._executor.all_nodes: if not self._executor.evaluation_context.is_done(applied_ptransform): pending_bundles = self._executor.node_to_pending_bundles.get( applied_ptransform, []) for bundle in pending_bundles: self._executor.schedule_consumption( applied_ptransform, bundle, [], self._executor.default_completion_callback) self._executor.node_to_pending_bundles[applied_ptransform] = []