diff --git a/CHANGES.md b/CHANGES.md index 537c3a32447c..ba38855228af 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -70,6 +70,7 @@ ## New Features / Improvements +* (Python) Reduced runtime type-check overhead by combining input and output checks into a single wrapper ([#20226](https://github.com/apache/beam/issues/20226)). * (Java) Flink runner: sources now forward Beam metrics to Flink at most once per second instead of after every record, which was a significant share of the CPU of reading. Metrics in the pipeline result are not affected ([#40607](https://github.com/apache/beam/issues/40607)). * (Go) Added `wait.On`, which delays each input window until the corresponding windows in its signal PCollections have closed ([#39909](https://github.com/apache/beam/issues/39909)). * (Python) Expanded the SDK worker heap dump (`--experiments=enable_heap_dump`) with process RSS, CPython allocator/GC stats, and glibc `mallinfo2` native-heap/fragmentation stats to help distinguish native-heap from Python-object memory growth ([#39244](https://github.com/apache/beam/issues/39244)). diff --git a/sdks/python/apache_beam/typehints/typecheck.py b/sdks/python/apache_beam/typehints/typecheck.py index 7e84779a0f15..d8818298de0f 100644 --- a/sdks/python/apache_beam/typehints/typecheck.py +++ b/sdks/python/apache_beam/typehints/typecheck.py @@ -84,25 +84,67 @@ def teardown(self): return self.dofn.teardown() -class OutputCheckWrapperDoFn(AbstractDoFnWrapper): - """A DoFn that verifies against common errors in the output type.""" - def __init__(self, dofn, full_label): +class RuntimeTypeCheckWrapperDoFn(AbstractDoFnWrapper): + """A wrapper around a DoFn which performs type-checking of input and output.""" + def __init__(self, dofn, type_hints, full_label): super().__init__(dofn) self.full_label = full_label + process_fn = self.dofn._process_argspec_fn() + if hasattr(process_fn, '__func__'): + self._process_fn = process_fn.__func__ + self._process_fn_needs_self = True + else: + self._process_fn = process_fn + self._process_fn_needs_self = False + + if type_hints.input_types: + input_args, input_kwargs = type_hints.input_types + self._input_hints = getcallargs_forhints( + process_fn, *input_args, **input_kwargs) + else: + self._input_hints = None + + # TODO(robertwb): Multi-output. + self._output_type_hint = type_hints.simple_output_type(full_label) + def wrapper(self, method, args, kwargs): try: - result = method(*args, **kwargs) + result = self._type_check_result(method(*args, **kwargs)) except TypeCheckError as e: - # TODO(BEAM-10710): Remove the 'ParDo' prefix for the label name - error_msg = ( - 'Runtime type violation detected within ParDo(%s): ' - '%s' % (self.full_label, e)) - _, _, tb = sys.exc_info() - raise TypeCheckError(error_msg).with_traceback(tb) + raise self._add_label(e) else: return self._check_type(result) + def process(self, *args, **kwargs): + try: + if self._input_hints: + if self._process_fn_needs_self: + actual_inputs = inspect.getcallargs( + self._process_fn, self.dofn, *args, **kwargs) # pylint: disable=deprecated-method + else: + actual_inputs = inspect.getcallargs(self._process_fn, *args, **kwargs) # pylint: disable=deprecated-method + + for var, hint in self._input_hints.items(): + if hint is actual_inputs[var]: + # self parameter + continue + _check_instance_type(hint, actual_inputs[var], var, True) + + result = self._type_check_result(self.dofn.process(*args, **kwargs)) + except TypeCheckError as e: + raise self._add_label(e) + else: + return self._check_type(result) + + def _add_label(self, e): + # TODO(BEAM-10710): Remove the 'ParDo' prefix for the label name. + error_msg = ( + 'Runtime type violation detected within ParDo(%s): ' + '%s' % (self.full_label, e)) + _, _, tb = sys.exc_info() + return TypeCheckError(error_msg).with_traceback(tb) + @staticmethod def _check_type(output): if output is None: @@ -120,36 +162,6 @@ def _check_type(output): 'iterable. %s was returned instead.' % type(output)) return output - -class TypeCheckWrapperDoFn(AbstractDoFnWrapper): - """A wrapper around a DoFn which performs type-checking of input and output. - """ - def __init__(self, dofn, type_hints, label=None): - super().__init__(dofn) - self._process_fn = self.dofn._process_argspec_fn() - if type_hints.input_types: - input_args, input_kwargs = type_hints.input_types - self._input_hints = getcallargs_forhints( - self._process_fn, *input_args, **input_kwargs) - else: - self._input_hints = None - # TODO(robertwb): Multi-output. - self._output_type_hint = type_hints.simple_output_type(label) - - def wrapper(self, method, args, kwargs): - result = method(*args, **kwargs) - return self._type_check_result(result) - - def process(self, *args, **kwargs): - if self._input_hints: - actual_inputs = inspect.getcallargs(self._process_fn, *args, **kwargs) # pylint: disable=deprecated-method - for var, hint in self._input_hints.items(): - if hint is actual_inputs[var]: - # self parameter - continue - _check_instance_type(hint, actual_inputs[var], var, True) - return self._type_check_result(self.dofn.process(*args, **kwargs)) - def _type_check_result(self, transform_results): if self._output_type_hint is None or transform_results is None: return transform_results @@ -289,11 +301,9 @@ def visit_transform(self, applied_transform): if isinstance(transform.fn, core.CombineValuesDoFn): transform.fn.combinefn = self._wrapped_fn else: - transform.fn = transform.dofn = OutputCheckWrapperDoFn( - TypeCheckWrapperDoFn( - transform.fn, - transform.get_type_hints(), - applied_transform.full_label), + transform.fn = transform.dofn = RuntimeTypeCheckWrapperDoFn( + transform.fn, + transform.get_type_hints(), applied_transform.full_label) diff --git a/sdks/python/apache_beam/typehints/typecheck_test.py b/sdks/python/apache_beam/typehints/typecheck_test.py index 81046cd3cc0f..16f7d976046f 100644 --- a/sdks/python/apache_beam/typehints/typecheck_test.py +++ b/sdks/python/apache_beam/typehints/typecheck_test.py @@ -34,10 +34,14 @@ from apache_beam import Pipeline from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import TypeOptions +from apache_beam.pipeline import PipelineVisitor from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to +from apache_beam.transforms import combiners +from apache_beam.transforms import userstate from apache_beam.typehints import decorators +from apache_beam.typehints import typecheck from apache_beam.typehints import with_input_types from apache_beam.typehints import with_output_types @@ -300,5 +304,90 @@ def process(self, element, *args, **kwargs): self.p.run().wait_until_finish() +@with_input_types(int) +@with_output_types(int) +class _AddOneDoFn(beam.DoFn): + def process(self, element): + yield element + 1 + + +class RuntimeTypeCheckWrapperDoFnTest(unittest.TestCase): + def test_visitor_applies_single_wrapper_layer(self): + p = beam.Pipeline(options=PipelineOptions(runtime_type_check=True)) + _ = (p | beam.Create([1, 2]) | 'TypedStep' >> beam.ParDo(_AddOneDoFn())) + + p.visit(typecheck.TypeCheckVisitor()) + + wrapped_dofns = [] + + class _Collector(PipelineVisitor): + def visit_transform(self, applied_transform): + transform = applied_transform.transform + if isinstance(transform, beam.ParDo) and isinstance( + getattr(transform, 'fn', None), typecheck.AbstractDoFnWrapper): + wrapped_dofns.append(transform.fn) + + p.visit(_Collector()) + + self.assertTrue(wrapped_dofns) + for wrapper in wrapped_dofns: + self.assertIsInstance(wrapper, typecheck.RuntimeTypeCheckWrapperDoFn) + self.assertNotIsInstance(wrapper.dofn, typecheck.AbstractDoFnWrapper) + + def test_wrapper_labels_type_check_errors(self): + dofn = _AddOneDoFn() + wrapper = typecheck.RuntimeTypeCheckWrapperDoFn( + dofn, dofn.get_type_hints(), 'MyStep') + + with self.assertRaisesRegex( + typecheck.TypeCheckError, + r'Runtime type violation detected within ParDo\(MyStep\)'): + wrapper.process('not-an-int') + + def test_wrapper_preserves_results(self): + dofn = _AddOneDoFn() + wrapper = typecheck.RuntimeTypeCheckWrapperDoFn( + dofn, dofn.get_type_hints(), 'MyStep') + + self.assertEqual(list(wrapper.process(1)), [2]) + + +def _make_stateful_dofn(): + count_state = userstate.CombiningValueStateSpec( + 'count', combiners.CountCombineFn()) + + class _CountingStatefulDoFn(beam.DoFn): + def process(self, element, count=beam.DoFn.StateParam(count_state)): + count.add(1) + yield element[1] + + return _CountingStatefulDoFn() + + +class RuntimeTypeCheckStatefulDoFnTest(unittest.TestCase): + def test_wrapper_does_not_cache_bound_process(self): + dofn = _make_stateful_dofn() + wrapper = typecheck.RuntimeTypeCheckWrapperDoFn( + dofn, dofn.get_type_hints(), 'Step') + + for attr, value in wrapper.__dict__.items(): + self.assertFalse( + hasattr(value, '__self__'), 'wrapper caches bound method %r' % attr) + + userstate.validate_stateful_dofn(wrapper) + + def test_stateful_dofn_with_runtime_type_check(self): + options = PipelineOptions() + options.view_as(TypeOptions).runtime_type_check = True + + with TestPipeline(options=options) as p: + result = ( + p + | beam.Create([('k', 1), ('k', 2), ('k', 3)]) + | beam.ParDo(_make_stateful_dofn())) + + assert_that(result, equal_to([1, 2, 3])) + + if __name__ == '__main__': unittest.main()