diff --git a/sdks/python/apache_beam/runners/worker/bundle_processor.py b/sdks/python/apache_beam/runners/worker/bundle_processor.py index ad48358d588e..85f1e43d6039 100644 --- a/sdks/python/apache_beam/runners/worker/bundle_processor.py +++ b/sdks/python/apache_beam/runners/worker/bundle_processor.py @@ -24,6 +24,7 @@ import base64 import bisect import collections +import concurrent.futures import copy import heapq import itertools @@ -76,6 +77,7 @@ from apache_beam.runners.worker import operation_specs from apache_beam.runners.worker import operations from apache_beam.runners.worker import statesampler +from apache_beam.runners.worker.worker_status import thread_dump from apache_beam.transforms import TimeDomain from apache_beam.transforms import core from apache_beam.transforms import environments @@ -1130,7 +1132,30 @@ def __init__( 'fnapi-step-%s' % self.process_bundle_descriptor.id, self.counter_factory) - self.ops = self.create_execution_tree(self.process_bundle_descriptor) + with concurrent.futures.ThreadPoolExecutor( + max_workers=1, thread_name_prefix='ExecutionTreeCreator') as executor: + future = executor.submit( + self.create_execution_tree, self.process_bundle_descriptor) + try: + self.ops = future.result(timeout=3600) + except concurrent.futures.TimeoutError: + # In rare cases, unpickling a DoFn might get permanently stuck, + # for example when unpickling involves importing a module and + # a subprocess is launched during the import operation. + _LOGGER.error( + 'Timed out while reconstructing a pipeline fragment for: %s.\n' + 'This is likely a transient error. The SDK harness ' + 'will self-terminate, and the runner can retry the operation. ' + 'If the error is frequent, check whether the stuckness happens ' + 'while deserializing (unpickling) a dependency of your pipeline ' + 'in the stacktrace below: \n%s\n', + self.process_bundle_descriptor.id, + thread_dump('ExecutionTreeCreator')) + # Raising an exception here doesn't interrupt the left-over thread. + # Out of caution, terminate the SDK harness process. + from apache_beam.runners.worker.sdk_worker_main import terminate_sdk_harness + terminate_sdk_harness() + for op in reversed(self.ops.values()): op.setup(self.data_sampler) self.splitting_lock = threading.Lock() diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker_main.py b/sdks/python/apache_beam/runners/worker/sdk_worker_main.py index 7ea0e0eb1099..cdb807e8dbc5 100644 --- a/sdks/python/apache_beam/runners/worker/sdk_worker_main.py +++ b/sdks/python/apache_beam/runners/worker/sdk_worker_main.py @@ -233,6 +233,10 @@ def terminate_sdk_harness(): if _FN_LOG_HANDLER: _FN_LOG_HANDLER.close() os.kill(os.getpid(), signal.SIGINT) + # Delay further control flow in the caller until process is terminated. + time.sleep(60) + # Try to force-terminate if still running. + os.kill(os.getpid(), signal.SIGKILL) def _load_pipeline_options(options_json): diff --git a/sdks/python/apache_beam/runners/worker/worker_status.py b/sdks/python/apache_beam/runners/worker/worker_status.py index f4102b193895..1d54a3ee1764 100644 --- a/sdks/python/apache_beam/runners/worker/worker_status.py +++ b/sdks/python/apache_beam/runners/worker/worker_status.py @@ -66,13 +66,23 @@ def _current_frames(): return sys._current_frames() # pylint: disable=protected-access -def thread_dump(): - """Get a thread dump for the current SDK worker harness. """ +def thread_dump(thread_prefix=None): + """Get a thread dump for the current SDK harness. + + Args: + thread_prefix: (str) An optional prefix to filter threads by. + """ # deduplicate threads with same stack trace stack_traces = defaultdict(list) frames = _current_frames() - for t in threading.enumerate(): + threads_to_dump = threading.enumerate() + if thread_prefix: + threads_to_dump = [ + t for t in threads_to_dump if t.name.startswith(thread_prefix) + ] + + for t in threads_to_dump: try: stack_trace = ''.join(traceback.format_stack(frames[t.ident])) except KeyError: