#
# 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 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] = []